Update app.py

app.py

@@ -1,17 +1,13 @@
-from transformers import pipeline, Conversation
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM, TextGenerationPipeline
 import gradio as gr
 
-chatbot = pipeline(model="Writer/palmyra-small")
+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)
 
-message_list = []
-response_list = []
+text_generator = TextGenerationPipeline(model=model, tokenizer=tokenizer)
 
-def vanilla_chatbot(message, history):
-    conversation = Conversation(text=message, past_user_inputs=message_list, generated_responses=response_list)
-    conversation = chatbot(conversation)
-
-    return conversation.generated_responses[-1]
-
-demo_chatbot = gr.ChatInterface(vanilla_chatbot, title="Writer - Palmyra-small Chatbot", description="Enter text to start chatting.")
-
-demo_chatbot.launch()
+iface = gr.Interface(fn=text_generator, inputs="text", outputs="text")
+iface.launch()