Update app.py

app.py

@@ -1,57 +1,24 @@
-import os
 import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer
-import torch
 
-# Load model and tokenizer
-def load_model():
-    model_name = "Viet-Mistral/Vistral-7B-Chat"
-    # Lấy token từ environment variable
+# Hàm để sử dụng mô hình trả lời câu hỏi
+def chatbot(input_text):
+    # Đảm bảo bạn đã load mô hình và tokenizer
+    model_name = "Viet-Mistral/Vistral-7B-Chat"  # Hoặc tên mô hình bạn đang sử dụng
     hf_token = os.getenv("HF_AUTH_TOKEN")
     model = AutoModelForCausalLM.from_pretrained(model_name, token=hf_token)
     tokenizer = AutoTokenizer.from_pretrained(model_name, token=hf_token)
-    return model, tokenizer
-
-# Setup and load the model
-model, tokenizer = load_model()
-
-# Generate response based on conversation history
-def generate(messages):
-    prompt_text = ""
-    for message in messages:
-        role = message["role"]
-        content = message["content"]
-        if role == "user":
-            prompt_text += f"User: {content}\n"
-        else:
-            prompt_text += f"Assistant: {content}\n"
-    prompt_text += "Assistant: "  # để chuẩn bị cho model generate tiếp
-
-    # Tokenize input prompt
-    inputs = tokenizer(prompt_text, return_tensors="pt")
     
-    # Generate response
-    with torch.no_grad():
-        output = model.generate(inputs.input_ids, max_length=512, num_return_sequences=1, pad_token_id=tokenizer.eos_token_id)
-
-    response = tokenizer.decode(output[0], skip_special_tokens=True)
+    # Xử lý đầu vào và tạo ra câu trả lời
+    inputs = tokenizer(input_text, return_tensors="pt")
+    outputs = model.generate(inputs["input_ids"])
+    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
     
-    return response.strip()
-
-# Gradio interface
-def chatbot_interface():
-    with gr.Blocks() as demo:
-        gr.Markdown("# Chatbot sử dụng Vistral của Viet-Mistral")
-        chatbox = gr.Chatbot()
-        message = gr.Textbox(placeholder="Gửi tin nhắn...")
-        send_button = gr.Button("Gửi")
-
-        send_button.click(generate, inputs=message, outputs=chatbox)
-
-    return demo
+    # Trả về câu trả lời
+    return response
 
-# Main function to run the app
-if __name__ == "__main__":
-    demo = chatbot_interface()
-    demo.launch(share=True)
+# Tạo giao diện Gradio
+iface = gr.Interface(fn=chatbot, inputs="text", outputs="text", live=True)
 
+# Khởi động giao diện
+iface.launch()