Update app.py

app.py

@@ -1,10 +1,8 @@
 # app.py
 import torch
 import torch.nn as nn
-from fastapi import FastAPI, Request
-from fastapi.responses import JSONResponse
 from tokenizers import Tokenizer
-import json
+import gradio as gr
 
 # -----------------------------
 # Load tokenizer
@@ -14,7 +12,7 @@ tokenizer = Tokenizer.from_file(tokenizer_path)
 vocab_size = tokenizer.get_vocab_size()
 
 # -----------------------------
-# Define your same transformer structure
+# Define the same transformer as used in training
 # -----------------------------
 class SimpleTransformer(nn.Module):
     def __init__(self, vocab_size, d_model=128, nhead=4, num_layers=4):
@@ -39,25 +37,28 @@ model.load_state_dict(torch.load("pytorch_model.bin", map_location="cpu"))
 model.eval()
 
 # -----------------------------
-# FastAPI app
+# Generation function
 # -----------------------------
-app = FastAPI()
-
-@app.post("/generate")
-async def generate(request: Request):
-    data = await request.json()
-    prompt = data.get("prompt", "")
-
-    # Tokenize
+def generate(prompt):
     input_ids = tokenizer.encode(prompt).ids
     input_tensor = torch.tensor([input_ids])
-
-    # Generate (basic greedy)
+    
     with torch.no_grad():
         output = model(input_tensor)
         predicted_ids = torch.argmax(output, dim=-1)[0].tolist()
-
-    # Decode
+    
     response_text = tokenizer.decode(predicted_ids)
+    return response_text
+
+# -----------------------------
+# Gradio Interface
+# -----------------------------
+iface = gr.Interface(
+    fn=generate,
+    inputs=gr.Textbox(lines=2, placeholder="Type a prompt for Dave..."),
+    outputs="text",
+    title="Dave – Fully Custom AI",
+    description="Interact with your fully custom AI trained from scratch."
+)
 
-    return JSONResponse({"response": response_text})
+iface.launch()