Update app.py

app.py

@@ -1,15 +1,63 @@
-# app.py
-from transformers import GPT2Tokenizer, GPT2LMHeadModel
 import gradio as gr
+from transformers import GPT2Config, GPT2LMHeadModel, GPT2Tokenizer
+from datasets import load_dataset
+import torch
 
-# Load the tiny pretrained GPT model
-tokenizer = GPT2Tokenizer.from_pretrained("./tiny-gpt")
-model = GPT2LMHeadModel.from_pretrained("./tiny-gpt")
+# Load tokenizer
+tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
+tokenizer.pad_token = tokenizer.eos_token
 
-def generate_text(prompt):
-    inputs = tokenizer(prompt, return_tensors="pt")
-    outputs = model.generate(**inputs, max_length=50)
+# EXTREMELY tiny model
+config = GPT2Config(
+    vocab_size=tokenizer.vocab_size,
+    n_embd=64,
+    n_layer=1,
+    n_head=1,
+    n_positions=64,
+    n_ctx=64
+)
+
+model = GPT2LMHeadModel(config)
+
+# Load VERY small dataset
+dataset = load_dataset("wikitext", "wikitext-2-raw-v1", split="train[:0.1%]")
+
+def tokenize(x):
+    return tokenizer(
+        x["text"],
+        truncation=True,
+        padding="max_length",
+        max_length=64
+    )
+
+dataset = dataset.map(tokenize)
+dataset.set_format(type="torch", columns=["input_ids"])
+
+# ONE tiny training loop
+optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)
+
+for step, batch in enumerate(dataset):
+    if step > 50:  # HARD LIMIT
+        break
+    outputs = model(
+        input_ids=batch["input_ids"].unsqueeze(0),
+        labels=batch["input_ids"].unsqueeze(0)
+    )
+    loss = outputs.loss
+    loss.backward()
+    optimizer.step()
+    optimizer.zero_grad()
+
+print("Tiny training finished")
+
+def generate(text):
+    inputs = tokenizer(text, return_tensors="pt")
+    outputs = model.generate(inputs["input_ids"], max_length=50)
     return tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-# Gradio web interface
-gr.Interface(fn=generate_text, inputs="text", outputs="text").launch()
+gr.Interface(
+    fn=generate,
+    inputs="text",
+    outputs="text",
+    title="agent_1 (trained inside Space)"
+).launch()