Update app.py

app.py

@@ -1,10 +1,24 @@
 import gradio as gr
-from transformers import pipeline
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from peft import PeftModel
+
+base_model_name = "nltpt/Llama-3.2-1B-Instruct"
+base_model = AutoModelForCausalLM.from_pretrained(base_model_name)
+
+tokenizer = AutoTokenizer.from_pretrained(base_model_name)
+
+peft_model_path = "genaforvena/huivam_finnegan_llama3.2-1b"
+model = PeftModel.from_pretrained(base_model, peft_model_path)
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+model.to(device)
 
 def reply(prompt):
-    generator = pipeline("text-generation", model="genaforvena/huivam_finnegan_llama3.2-1b")
-    output = generator(prompt, max_new_tokens=150)
-    return output[0]['generated_text']
+    input_ids = tokenizer.encode(prompt, return_tensors="pt")
+    output = model.generate(input_ids, max_new_tokens=250, num_return_sequences=1)
+    generated_text = tokenizer.decode(output[0], skip_special_tokens=True)
+
+    return generated_text
 
 demo = gr.Interface(fn=reply, inputs="text", outputs="text")
 demo.launch()