Update app.py

app.py

@@ -1,38 +1,19 @@
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
 import torch
 
-# Load model and tokenizer
-model_name = "mistralai/Mistral-7B-Instruct-v0.1"
+# Load free public model
+model_name = "google/flan-t5-small"
 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 chat_with_mistral(user_input):
-    prompt = f"<s>[INST] {user_input.strip()} [/INST]"
-    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
-    output = model.generate(
-        **inputs,
-        max_new_tokens=256,
-        temperature=0.7,
-        top_p=0.9,
-        do_sample=True
-    )
-    response = tokenizer.decode(output[0], skip_special_tokens=True)
-    # Remove the prompt portion
-    if "[/INST]" in response:
-        response = response.split("[/INST]")[1].strip()
-    return response
+def chat(prompt):
+    input_ids = tokenizer(prompt, return_tensors="pt").input_ids
+    with torch.no_grad():
+        outputs = model.generate(input_ids, max_new_tokens=200)
+    reply = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return reply
 
-iface = gr.Interface(
-    fn=chat_with_mistral,
-    inputs=gr.Textbox(lines=2, placeholder="Ask something..."),
-    outputs=gr.Textbox(),
-    title="Mistral 7B Chatbot",
-    description="A chatbot powered by Mistral-7B-Instruct-v0.1."
-)
+demo = gr.Interface(fn=chat, inputs="text", outputs="text", title="FLAN-T5 Chatbot")
 
-iface.launch()
+demo.launch()