Update app.py

app.py

@@ -1,20 +1,14 @@
-from peft import PeftModel
-from transformers import LLaMATokenizer, LLaMAForCausalLM, GenerationConfig
-import torch
-n_gpus = torch.cuda.device_count()
-max_memory = {i: max_memory for i in range(n_gpus)}
 
-print(f'Max memory : {max_memory}')
+from peft import PeftModel
+from transformers import LlamaTokenizer, LlamaForCausalLM, GenerationConfig
 
-tokenizer = LLaMATokenizer.from_pretrained("decapoda-research/llama-7b-hf")
-max_memory = '40GB'
 
-model = LLaMAForCausalLM.from_pretrained(
+tokenizer = LlamaTokenizer.from_pretrained("decapoda-research/llama-7b-hf")
+model = LlamaForCausalLM.from_pretrained(
     "decapoda-research/llama-7b-hf",
     load_in_8bit=True,
-    device_map="auto",max_memory=max_memory
+    device_map="auto",
 )
-
 model = PeftModel.from_pretrained(model, "tloen/alpaca-lora-7b")
 
 def generate_prompt(instruction, input=None):
@@ -57,11 +51,24 @@ def evaluate(instruction, input=None):
         output = tokenizer.decode(s)
         print("Response:", output.split("### Response:")[1].strip())
 
-import gradio as gr
+import streamlit as st
 from peft import PeftModel
-from transformers import LLaMATokenizer, LLaMAForCausalLM, GenerationConfig
+from transformers import LlamaTokenizer, LlamaForCausalLM, GenerationConfig
+
+model_name = 'bhaskar/LLaMA-7B-peft'
+tokenizer = LlamaTokenizer.from_pretrained(model_name)
+model = LlamaForCausalLM.from_pretrained(model_name).cuda()
+generation_config = GenerationConfig(
+    do_sample=True,
+    max_length=1024,
+    top_p=0.9,
+    temperature=1.0,
+    no_repeat_ngram_size=3,
+    num_return_sequences=1,
+)
 
-import gradio as gr
+def generate_prompt(instruction):
+    return f"### Instruction: {instruction}\n\n### Response:"
 
 def evaluate1(instruction):
     prompt = generate_prompt(instruction)
@@ -78,11 +85,15 @@ def evaluate1(instruction):
         output = tokenizer.decode(s)
         return output.split("### Response:")[1].strip()
 
-inputs = gr.inputs.Textbox(lines=5, label="Instruction")
-outputs = gr.outputs.Textbox(label="Response")
-title = "LLaMA-7B Language Model"
-description = "This is a LLaMA-7B language model fine-tuned on various text datasets to generate text for a given task. It was trained on PyTorch by and is capable of generating high-quality, coherent text that is similar to human writing. The model is highly versatile and can be used for a variety of tasks, including text completion, summarization, and translation."
-copyright = "Copyright Bhaskar Tripathi (2023)"
-
-gr.Interface(evaluate1, inputs, outputs, title=title, description=description, footer=copyright, flag=False).launch()
-
+def main():
+    st.set_page_config(page_title="LLaMA-7B Language Model")
+    st.title("LLaMA-7B Language Model")
+    st.write("This is a LLaMA-7B language model fine-tuned on various text datasets to generate text for a given task. It was trained on PyTorch by and is capable of generating high-quality, coherent text that is similar to human writing. The model is highly versatile and can be used for a variety of tasks, including text completion, summarization, and translation.")
+    instruction = st.text_area("Instruction", height=200)
+    if st.button("Generate Response"):
+        with st.spinner("Generating response..."):
+            output = evaluate1(instruction)
+        st.write(output)
+
+if __name__ == "__main__":
+    main()