Update app.py

app.py

@@ -1,3 +1,4 @@
+import gradio as gr
 import torch
 from peft import PeftModel, PeftConfig
 from transformers import AutoModelForCausalLM, AutoTokenizer, GenerationConfig
@@ -41,7 +42,6 @@ class Conversation:
         final_text += DEFAULT_RESPONSE_TEMPLATE
         return final_text.strip()
 
-
 def generate(model, tokenizer, prompt, generation_config):
     data = tokenizer(prompt, return_tensors="pt", add_special_tokens=False)
     data = {k: v.to(model.device) for k, v in data.items()}
@@ -69,17 +69,22 @@ model.eval()
 
 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, use_fast=False)
 generation_config = GenerationConfig.from_pretrained(MODEL_NAME)
-print(generation_config)
 
-inputs = ["Почему трава зеленая?", "Сочини длинный рассказ, обязательно упоминая следующие объекты. Дано: Таня, мяч"]
-for inp in inputs:
+# Gradio interface setup
+def chat_with_model(user_input):
     conversation = Conversation()
-    conversation.add_user_message(inp)
+    conversation.add_user_message(user_input)
     prompt = conversation.get_prompt(tokenizer)
+    response = generate(model, tokenizer, prompt, generation_config)
+    conversation.add_bot_message(response)
+    return conversation.messages[-1]["content"]
+
+iface = gr.Interface(
+    fn=chat_with_model,
+    inputs=gr.Textbox(prompt="You:"),
+    outputs=gr.Textbox(prompt="Bot:"),
+    live=True,
+    title="Chat with Bot",
+)
 
-    output = generate(model, tokenizer, prompt, generation_config)
-    print(inp)
-    print(output)
-    print()
-    print("==============================")
-    print()
+iface.launch()