Update app.py

app.py

@@ -1,43 +1,55 @@
 import os
-import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer, TextStreamer
-import spaces
 import gradio as gr
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
 
-@spaces.GPU
+# Load model and tokenizer
 def load_model():
-    model_id = "microsoft/phi-2"
-    access_token = os.environ.get("HF_AUTH_TOKEN")
+    model_name = "viet-ai/vistral-7b-chat"  # Vistral của Viet-Mistral
+    tokenizer = AutoTokenizer.from_pretrained(model_name, use_auth_token=os.getenv("HF_AUTH_TOKEN"))
+    model = AutoModelForCausalLM.from_pretrained(model_name, use_auth_token=os.getenv("HF_AUTH_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)
     
-    tokenizer = AutoTokenizer.from_pretrained(model_id, token=access_token)
-    model = AutoModelForCausalLM.from_pretrained(
-        model_id,
-        device_map="auto",
-        torch_dtype=torch.float16,
-        use_auth_token=access_token
-    )
-    streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
-    return model, tokenizer, streamer
-
-model, tokenizer, streamer = load_model()
-
-def generate(prompt, history):
-    messages = [{"role": "user", "content": prompt}]
-    prompt_text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
-
-    inputs = tokenizer(prompt_text, return_tensors="pt").to(model.device)
-    output = model.generate(
-        **inputs,
-        max_new_tokens=512,
-        do_sample=True,
-        temperature=0.8,
-        top_p=0.95,
-        top_k=50,
-        streamer=streamer
-    )
-    decoded_output = tokenizer.decode(output[0], skip_special_tokens=True)
-    # Tách phần phản hồi ra khỏi prompt
-    response = decoded_output.split(prompt_text)[-1].strip()
-    return response
-
-gr.ChatInterface(generate, title="💬 Chatbot Phi-2").launch()
+    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
+
+# Main function to run the app
+if __name__ == "__main__":
+    demo = chatbot_interface()
+    demo.launch(share=True)
+