Update app.py

app.py

@@ -1,5 +1,5 @@
 import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM, TextGenerationPipeline
+from transformers import AutoTokenizer, AutoModelForCausalLM
 import gradio as gr
 
 model_name = "Writer/palmyra-small"
@@ -7,7 +7,28 @@ tokenizer = AutoTokenizer.from_pretrained(model_name)
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 model = AutoModelForCausalLM.from_pretrained(model_name).to(device)
 
-text_generator = TextGenerationPipeline(model=model, tokenizer=tokenizer)
+def generate_response(prompt):
+    input_text_template = (
+        "A chat between a curious user and an artificial intelligence assistant. "
+        "The assistant gives helpful, detailed, and polite answers to the user's questions. "
+        f"USER: {prompt} "
+        "ASSISTANT:"
+    )
 
-iface = gr.Interface(fn=text_generator, inputs="text", outputs="text")
-iface.launch()
+    model_inputs = tokenizer(input_text_template, return_tensors="pt").to(device)
+
+    gen_conf = {
+        "top_k": 20,
+        "max_length": 200,
+        "temperature": 0.6,
+        "do_sample": True,
+        "eos_token_id": tokenizer.eos_token_id,
+    }
+
+    output = model.generate(**model_inputs, **gen_conf)
+
+    generated_text = tokenizer.decode(output[0], skip_special_tokens=True)
+    return generated_text
+
+iface = gr.Interface(fn=generate_response, inputs="text", outputs="text")
+iface.launch()