Update app.py

app.py

@@ -1,23 +1,28 @@
 import gradio as gr
 import torch
-from transformers import GPT2Tokenizer
-from evo_model import EvoDecoderModel
-from generate import generate_text
+from transformers import AutoTokenizer
+from generate import load_model, generate_text
 
-# Load tokenizer and model
-tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
+# Settings
+MODEL_PATH = "evo_decoder.pt"
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+TOKENIZER_NAME = "gpt2"  # Or your custom tokenizer
+MAX_TOKENS = 50
+
+# Load tokenizer & model
+tokenizer = AutoTokenizer.from_pretrained(TOKENIZER_NAME)
 vocab_size = tokenizer.vocab_size
+model = load_model(vocab_size, MODEL_PATH, DEVICE)
 
-model = EvoDecoderModel(vocab_size)
-model.load_state_dict(torch.load("evo_decoder.pt", map_location="cpu"))
-model.eval()
+def chat(prompt):
+    output = generate_text(model, tokenizer, prompt, max_new_tokens=MAX_TOKENS, device=DEVICE)
+    return output
 
-def chat_with_evo(prompt):
-    response = generate_text(model, tokenizer, prompt)
-    return response
+# Gradio UI
+iface = gr.Interface(fn=chat,
+                     inputs=gr.Textbox(label="Ask Evo..."),
+                     outputs=gr.Textbox(label="Evo's Response"),
+                     title="🧠 EvoDecoder Chatbot",
+                     description="A lightweight conversational model powered by EvoDecoder")
 
-gr.Interface(fn=chat_with_evo, 
-             inputs="text", 
-             outputs="text", 
-             title="🧠 EvoDecoder Chatbot",
-             description="Ask Evo anything. Powered by your trained EvoDecoder.").launch()
+iface.launch()