Update app.py

app.py

@@ -2,7 +2,8 @@ import gradio as gr
 import torch
 import torch.nn as nn
 from torch.nn import functional as F
-import json # <-- Added this import
+import json
+import os # <-- Added for file path checks
 
 # --- Model Definition (same as before) ---
 # NOTE: The model class MUST be defined in your app.py file
@@ -149,23 +150,35 @@ model.to(device)
 
 
 # --- Gradio UI & Inference function ---
-def generate_text(prompt, max_new_tokens):
+def generate_text_chat(message, history):
+    # We'll just use the most recent message as the prompt.
+    prompt = message
+    # You can adjust this to a different number of tokens if you like.
+    max_new_tokens = 50 
+
     # Encode the prompt text into tokens.
-    context = torch.tensor(encode(prompt), dtype=torch.long, device=device).unsqueeze(0)
+    encoded_prompt = [stoi.get(c, 0) for c in prompt]
+    if not encoded_prompt:
+        return "Prompt is empty or contains unknown characters."
+
+    context = torch.tensor(encoded_prompt, dtype=torch.long, device=device).unsqueeze(0)
     # Generate new tokens.
     generated_text_indices = model.generate(context, max_new_tokens=max_new_tokens)
     # Decode the tokens back into text.
-    return decode(generated_text_indices[0].tolist())
-
-demo = gr.Interface(
-    fn=generate_text,
-    inputs=[
-        gr.Textbox(label="Prompt", placeholder="Enter your text prompt here..."),
-        gr.Slider(1, 100, value=20, step=1, label="Number of new tokens to generate"),
-    ],
-    outputs="text",
-    title="Tiny Language Model",
-    description="A simple character-level language model trained in PyTorch."
+    generated_text = decode(generated_text_indices[0].tolist())
+    
+    # Return only the newly generated part of the text, removing the original prompt
+    return generated_text[len(prompt):]
+
+# Using gr.ChatInterface for a conversational experience
+demo = gr.ChatInterface(
+    fn=generate_text_chat,
+    title="Tiny Language Model Chat",
+    description="A simple character-level language model trained in PyTorch, now with a chat interface.",
+    # You can customize these components further if you like
+    chatbot=gr.Chatbot(height="500px"),
+    textbox=gr.Textbox(placeholder="Ask me anything...", container=False, scale=7),
+    theme="soft",
 )
 
-demo.launch()
+demo.launch()