Update app.py

app.py

@@ -1,36 +1,15 @@
- import gradio as gr
-from transformers import AutoTokenizer, AutoModelForCausalLM
-import torch
+import gradio as gr
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 
-MODEL_NAME = "mistralai/Mistral-7B-Instruct-v0.1"
+MODEL_NAME = "google/flan-t5-base"
 
 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-model = AutoModelForCausalLM.from_pretrained(
-    MODEL_NAME,
-    torch_dtype=torch.float16,
-    device_map="auto"
-)
+model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME)
 
 def generate_response(message, history):
-    prompt = f"""
-You are a professional AI assistant.
-Answer clearly and helpfully.
-
-User: {message}
-Assistant:
-"""
-
-    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
-
-    output = model.generate(
-        **inputs,
-        max_new_tokens=200,
-        temperature=0.7,
-        top_p=0.9
-    )
-
-    response = tokenizer.decode(output[0], skip_special_tokens=True)
-
-    return response.split("Assistant:")[-1].strip()
+    inputs = tokenizer(message, return_tensors="pt")
+    outputs = model.generate(**inputs, max_new_tokens=100)
+    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return response
 
 gr.ChatInterface(generate_response).launch()