Update app.py

app.py

@@ -6,6 +6,7 @@ from transformers import AutoTokenizer
 from datasets import Dataset
 import os
 
+# Tiny NanoGPT class
 class NanoGPT(nn.Module):
     def __init__(self, vocab_size=30522, n_embd=96, n_head=4, n_layer=3, block_size=96):
         super().__init__()
@@ -15,16 +16,15 @@ class NanoGPT(nn.Module):
         self.drop = nn.Dropout(0.1)
 
         self.layers = nn.ModuleList([
-            nn.TransformerDecoderLayer(
-                d_model=n_embd, nhead=n_head, dim_feedforward=n_embd*4,
-                dropout=0.1, activation="gelu", batch_first=True
-            ) for _ in range(n_layer)
+            nn.TransformerDecoderLayer(d_model=n_embd, nhead=n_head,
+                                       dim_feedforward=n_embd*4, dropout=0.1,
+                                       activation="gelu", batch_first=True)
+            for _ in range(n_layer)
         ])
 
         self.ln_f = nn.LayerNorm(n_embd)
         self.head = nn.Linear(n_embd, vocab_size, bias=False)
-        self.tok_emb.weight = self.head.weight  # weight tying
-        self.n_embd = n_embd
+        self.tok_emb.weight = self.head.weight  # tie weights
 
     def forward(self, idx, targets=None):
         B, T = idx.shape
@@ -33,16 +33,17 @@ class NanoGPT(nn.Module):
         x = self.drop(tok_emb + pos_emb)
 
         for layer in self.layers:
-            x = layer(x, None)  # causal self-attention
+            x = layer(x, None)  # self-attn only, causal
 
         x = self.ln_f(x)
         logits = self.head(x)
 
         if targets is None:
             return logits, None
+
         B, T, C = logits.shape
-        logits = logits.view(B*T, C)
-        targets = targets.view(B*T)
+        logits = logits.view(B * T, C)
+        targets = targets.view(B * T)
         loss = F.cross_entropy(logits, targets)
         return logits, loss
 
@@ -57,58 +58,62 @@ class NanoGPT(nn.Module):
             idx = torch.cat((idx, next_idx), dim=1)
         return idx
 
-# Globals
+# Setup
 tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
-vocab_size = tokenizer.vocab_size
-block_size = 96
-model = NanoGPT(vocab_size=vocab_size, n_embd=96, n_head=4, n_layer=3, block_size=block_size)
-
-model_path = "/data/nanogpt_yap.pt"  # /data is persistent on Spaces
+model = NanoGPT()
+model_path = "nanogpt_yap.pt"  # saved in current dir (non-persistent on restart, but ok for test)
 
 if os.path.exists(model_path):
     model.load_state_dict(torch.load(model_path, map_location="cpu"))
-    print("Loaded saved model")
+    print("Loaded existing model weights")
 
-# Tiny dataset (repeat for more tokens)
+# Small dataset for yapping about life (repeat for more data)
 life_texts = [
     "Life is what happens when you're busy making other plans.",
     "The meaning of life is to find your gift. The purpose is to give it away.",
     "You only live once, but if you do it right, once is enough.",
-    "Hey human, existence is weird. Coffee helps.",
-    "I think therefore I am... but mostly I just scroll.",
-    "Why do we exist? Probably for the memes and Java code.",
-    # add more if you want
+    "Existence is weird. Coffee helps sometimes.",
+    "I think therefore I am... mostly scrolling though.",
+    "Why do we exist? Probably for memes and Java bugs.",
+    "Another day, another existential crisis. Pass the tea.",
 ]
 
 def create_dataset():
-    text = " ".join(life_texts * 50)  # ~few k tokens
+    text = " ".join(life_texts * 50)  # small but repeated
     encodings = tokenizer(text, return_tensors="pt")
     input_ids = encodings.input_ids[0]
 
     seqs = []
-    for i in range(0, len(input_ids) - block_size - 1, block_size // 2):
-        chunk = input_ids[i:i + block_size + 1]
+    block_size = 96
+    step = block_size // 2
+    for i in range(0, len(input_ids) - block_size - 1, step):
+        chunk = input_ids[i : i + block_size + 1]
         if len(chunk) == block_size + 1:
             seqs.append(chunk)
 
     if not seqs:
         return None
-    data = {"input_ids": [s[:-1].tolist() for s in seqs], "labels": [s[1:].tolist() for s in seqs]}
+
+    data = {
+        "input_ids": [s[:-1].tolist() for s in seqs],
+        "labels": [s[1:].tolist() for s in seqs],
+    }
     return Dataset.from_dict(data)
 
-def train_once():
+def train_model():
     dataset = create_dataset()
-    if dataset is None:
-        return "Dataset too small!"
+    if dataset is None or len(dataset) == 0:
+        return "Dataset creation failed - too small!"
 
-    def collator(features):
+    def data_collator(features):
         batch = tokenizer.pad(features, padding=True, return_tensors="pt")
         batch["labels"] = batch["input_ids"].clone()
         return batch
 
     from transformers import Trainer, TrainingArguments
+
     args = TrainingArguments(
-        output_dir="/data/results",
+        output_dir="./results",
         num_train_epochs=5,
         per_device_train_batch_size=4,
         save_strategy="no",
@@ -116,51 +121,55 @@ def train_once():
         report_to="none",
         optim="adamw_torch",
         learning_rate=5e-4,
+        fp16=False,  # CPU
     )
 
     trainer = Trainer(
         model=model,
         args=args,
         train_dataset=dataset,
-        data_collator=collator,
+        data_collator=data_collator,
     )
 
     trainer.train()
     torch.save(model.state_dict(), model_path)
-    return "Training finished! Model saved to /data. Chat now!"
+    return "Training done! Model saved. You can chat now (responses may be silly)."
 
-def chat_with_nano(message, history):
-    if not message.strip():
-        return history + [["", "Say something existential... or about Java?"]]
+def generate_response(message, history):
+    if not message:
+        return history + [["", "Ask me something deep... or weird."]]
 
     prompt = f"Human: {message}\nAI: "
     inputs = tokenizer(prompt, return_tensors="pt").input_ids
 
     with torch.no_grad():
         generated = model.generate(inputs, max_new_tokens=80, temperature=0.95)
-    response = tokenizer.decode(generated[0][len(inputs[0]):], skip_special_tokens=True).strip()
+
+    full_text = tokenizer.decode(generated[0])
+    response = full_text[len(prompt):].strip()  # trim prompt part
 
     history.append([message, response])
     return history
 
+# Gradio UI
 with gr.Blocks() as demo:
-    gr.Markdown("# Nano Java/Life Yap AI")
-    gr.Markdown("Tiny ~1M param transformer. Train once, then chat!")
+    gr.Markdown("# Nano AI Yap Test")
+    gr.Markdown("Tiny from-scratch model (~1M params). Train first, then chat!")
 
     chatbot = gr.Chatbot(height=400)
-    msg = gr.Textbox(placeholder="Ask about life, existence, or Java...")
-    clear = gr.Button("Clear")
+    textbox = gr.Textbox(placeholder="Talk to me about life, existence, or anything...")
+    clear_btn = gr.Button("Clear Chat")
 
-    train_btn = gr.Button("Train Nano Model (10-60 min on CPU – do once!)")
-    status = gr.Textbox(label="Status")
+    train_button = gr.Button("Start Training (takes 10–60 min on free CPU – run once)")
+    status_box = gr.Textbox(label="Training Status", interactive=False)
 
-    train_btn.click(train_once, outputs=status)
+    train_button.click(train_model, outputs=status_box)
 
-    def respond(message, chat_history):
-        updated_history = chat_with_nano(message, chat_history)
-        return "", updated_history
+    def submit_chat(msg, hist):
+        updated_hist = generate_response(msg, hist)
+        return "", updated_hist
 
-    msg.submit(respond, [msg, chatbot], [msg, chatbot])
-    clear.click(lambda: None, None, chatbot, queue=False)
+    textbox.submit(submit_chat, [textbox, chatbot], [textbox, chatbot])
+    clear_btn.click(lambda: None, None, chatbot, queue=False)
 
 demo.launch(server_name="0.0.0.0", server_port=7860)