Update app.py

app.py

@@ -1,49 +1,74 @@
-# app.py – Hugging Face Space for Fully Custom "Dave" Model
 import torch
-from transformers import PreTrainedTokenizerFast, AutoModelForCausalLM
+import torch.nn as nn
+from tokenizers import Tokenizer
 from fastapi import FastAPI
 from pydantic import BaseModel
 import uvicorn
+import json
 
 # -----------------------------
-# Load tokenizer and model
+# Settings
 # -----------------------------
-tokenizer_path = "tokenizer.json"
-model_path = "pytorch_model.bin"
+MODEL_PATH = "./pytorch_model.bin"
+TOKENIZER_PATH = "./tokenizer.json"
+CONFIG_PATH = "./config.json"
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
 
-tokenizer = PreTrainedTokenizerFast(tokenizer_file=tokenizer_path)
-model = AutoModelForCausalLM.from_pretrained(
-    pretrained_model_name_or_path=".",
-    config="config.json",
-    state_dict=torch.load(model_path, map_location="cpu")
-)
+# -----------------------------
+# Load config
+# -----------------------------
+with open(CONFIG_PATH) as f:
+    config = json.load(f)
+
+# -----------------------------
+# Define the same architecture
+# -----------------------------
+class SimpleTransformer(nn.Module):
+    def __init__(self, vocab_size, d_model=128, nhead=4, num_layers=4):
+        super().__init__()
+        self.embedding = nn.Embedding(vocab_size, d_model)
+        encoder_layer = nn.TransformerEncoderLayer(d_model=d_model, nhead=nhead)
+        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
+        self.fc = nn.Linear(d_model, vocab_size)
 
+    def forward(self, x):
+        x = self.embedding(x)
+        x = x.transpose(0,1)
+        x = self.transformer(x)
+        x = x.transpose(0,1)
+        return self.fc(x)
+
+# -----------------------------
+# Load tokenizer and model
+# -----------------------------
+tokenizer = Tokenizer.from_file(TOKENIZER_PATH)
+vocab_size = config["vocab_size"]
+model = SimpleTransformer(vocab_size, config["d_model"], config["nhead"], config["num_layers"]).to(device)
+model.load_state_dict(torch.load(MODEL_PATH, map_location=device))
 model.eval()
 
 # -----------------------------
-# FastAPI app
+# FastAPI setup
 # -----------------------------
 app = FastAPI()
 
-class Prompt(BaseModel):
-    text: str
+class Query(BaseModel):
+    prompt: str
+    max_length: int = 64
 
 @app.post("/generate")
-def generate_text(prompt: Prompt):
-    inputs = tokenizer(prompt.text, return_tensors="pt")
+def generate(query: Query):
+    input_ids = tokenizer.encode(query.prompt).ids
+    input_tensor = torch.tensor([input_ids], device=device)
     with torch.no_grad():
-        outputs = model.generate(
-            **inputs,
-            max_length=64,
-            do_sample=True,
-            temperature=0.7,
-            top_p=0.9
-        )
-    decoded = tokenizer.batch_decode(outputs, skip_special_tokens=True)
-    return {"response": decoded[0]}
-
-# -----------------------------
-# Run the app (for local testing)
+        output = model(input_tensor)
+    predicted_ids = torch.argmax(output, dim=-1).squeeze().tolist()
+    response = tokenizer.decode(predicted_ids, skip_special_tokens=True)
+    return {"response": response}
+
+# -----------------------------
+# For running locally
 # -----------------------------
 if __name__ == "__main__":
     uvicorn.run(app, host="0.0.0.0", port=7860)