Update app.py

app.py

@@ -1,7 +1,13 @@
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM, TextGenerationPipeline
 import gradio as gr
 
-def greet(name):
-    return "Hello " + name + "!!"
+model_name = "Writer/palmyra-small"
+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)
 
-iface = gr.Interface(fn=greet, inputs="text", outputs="text")
+text_generator = TextGenerationPipeline(model=model, tokenizer=tokenizer)
+
+iface = gr.Interface(fn=text_generator, inputs="text", outputs="text")
 iface.launch()