Update app.py

app.py

@@ -1,34 +1,16 @@
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer, TextStreamer
 import gradio as gr
 
-model_id = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
+model_id = "google/flan-t5-small"
 
 tokenizer = AutoTokenizer.from_pretrained(model_id)
-model = AutoModelForCausalLM.from_pretrained(
-    model_id,
-    torch_dtype=torch.float32  # use float32 for CPU compatibility
-).to("cpu")
-
-streamer = TextStreamer(tokenizer, skip_prompt=True)
+model = AutoModelForSeq2SeqLM.from_pretrained(model_id).to("cpu")
 
 def chat(message, history):
-    prompt = ""
-    for user, bot in history:
-        prompt += f"<|user|>\n{user.strip()}\n<|assistant|>\n{bot.strip()}\n"
-    prompt += f"<|user|>\n{message.strip()}\n<|assistant|>\n"
-
+    prompt = f"Answer: {message}"
     inputs = tokenizer(prompt, return_tensors="pt").to("cpu")
-    outputs = model.generate(
-        **inputs,
-        max_new_tokens=256,
-        temperature=0.7,
-        top_p=0.9,
-        do_sample=True,
-        pad_token_id=tokenizer.eos_token_id
-    )
-    text = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    reply = text.split("<|assistant|>")[-1].strip()
-    return reply
+    outputs = model.generate(**inputs, max_new_tokens=100)
+    return tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-gr.ChatInterface(chat, title="TinyLlama Chat").launch()
+gr.ChatInterface(chat, title="Flan-T5 Chat (Small)").launch()