Update app.py

app.py

@@ -1,80 +1,82 @@
+# app.py
 import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM, Trainer, TrainingArguments
 from datasets import load_dataset
 import torch
 import os
 
-MODEL_DIR = "./model"
-
-def train_model():
-    model_name = "distilgpt2"
-
-    tokenizer = AutoTokenizer.from_pretrained(model_name)
-    model = AutoModelForCausalLM.from_pretrained(model_name)
-
-    dataset = load_dataset("text", data_files={"train": "data.txt"})
-
-    def tokenize(example):
-        return tokenizer(example["text"], truncation=True, padding="max_length", max_length=128)
-
-    tokenized = dataset.map(tokenize, batched=True)
-
-    training_args = TrainingArguments(
-        output_dir=MODEL_DIR,
-        per_device_train_batch_size=2,
-        num_train_epochs=2,
-        logging_steps=10,
-        save_steps=50
-    )
-
-    trainer = Trainer(
-        model=model,
-        args=training_args,
-        train_dataset=tokenized["train"],
+# -----------------------------
+# 1️⃣ Model setup
+# -----------------------------
+MODEL_DIR = "model"
+MODEL_NAME = "sshleifer/tiny-gpt2"  # tiny GPT-2, CPU-friendly
+
+# Load tokenizer & model
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
+
+# Fix padding issue
+tokenizer.pad_token = tokenizer.eos_token
+
+# -----------------------------
+# 2️⃣ Dataset setup
+# -----------------------------
+# Make sure you have 'data.txt' in the same folder as app.py
+dataset = load_dataset("text", data_files="data.txt")
+
+def tokenize(example):
+    return tokenizer(
+        example["text"],
+        truncation=True,
+        padding="max_length",
+        max_length=64  # small for CPU
     )
 
+tokenized_dataset = dataset.map(tokenize, batched=True)
+
+# -----------------------------
+# 3️⃣ Training setup
+# -----------------------------
+training_args = TrainingArguments(
+    output_dir=MODEL_DIR,
+    overwrite_output_dir=True,
+    per_device_train_batch_size=1,  # CPU-friendly
+    num_train_epochs=1,             # short test run
+    logging_steps=5,
+    save_steps=20,
+    save_total_limit=1
+)
+
+trainer = Trainer(
+    model=model,
+    args=training_args,
+    train_dataset=tokenized_dataset["train"]
+)
+
+# -----------------------------
+# 4️⃣ Gradio interface
+# -----------------------------
+def train_model():
     trainer.train()
+    return "✅ Training complete! Model saved to /model"
 
-    model.save_pretrained(MODEL_DIR)
-    tokenizer.save_pretrained(MODEL_DIR)
-
-    return "Training complete!"
-
-def load_model():
-    tokenizer = AutoTokenizer.from_pretrained(MODEL_DIR)
-    model = AutoModelForCausalLM.from_pretrained(MODEL_DIR)
-    return tokenizer, model
-
-def chat(user_input):
-    if not os.path.exists(MODEL_DIR):
-        return "Model not trained yet. Click Train first."
-
-    tokenizer, model = load_model()
-
-    prompt = f"User: {user_input}\nAssistant:"
-    inputs = tokenizer(prompt, return_tensors="pt")
-
-    outputs = model.generate(
-        **inputs,
-        max_length=100,
-        do_sample=True,
-        temperature=0.7
-    )
-
-    return tokenizer.decode(outputs[0], skip_special_tokens=True)
+def generate_text(prompt):
+    inputs = tokenizer(prompt, return_tensors="pt", padding=True)
+    output = model.generate(**inputs, max_length=64, pad_token_id=tokenizer.eos_token_id)
+    return tokenizer.decode(output[0], skip_special_tokens=True)
 
 with gr.Blocks() as demo:
-    gr.Markdown("# 🤖 My First AI")
-
-    train_btn = gr.Button("Train Model")
-    output = gr.Textbox()
-
-    train_btn.click(train_model, outputs=output)
-
-    user_input = gr.Textbox(label="Your message")
-    chat_output = gr.Textbox(label="AI Response")
-
-    send_btn = gr.Button("Send")
-    send_btn.click(chat, inputs=user_input, outputs=chat_output)
-
-demo.launch()
+    gr.Markdown("# Tiny AI Training Demo")
+    
+    with gr.Tab("Train Model"):
+        train_button = gr.Button("Train")
+        train_output = gr.Textbox(label="Logs")
+        train_button.click(train_model, outputs=train_output)
+    
+    with gr.Tab("Generate Text"):
+        prompt_input = gr.Textbox(label="Prompt")
+        generate_button = gr.Button("Generate")
+        generate_output = gr.Textbox(label="Output")
+        generate_button.click(generate_text, inputs=prompt_input, outputs=generate_output)
+
+demo.launch(share=True)