Update app.py

app.py

@@ -1,51 +1,121 @@
 import gradio as gr
-from datasets import load_dataset
-from transformers import AutoModelForCausalLM, AutoTokenizer, Trainer, TrainingArguments
+from transformers import (
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    Trainer,
+    TrainingArguments,
+    DataCollatorForLanguageModeling,
+)
+from datasets import load_dataset, Dataset
+import torch
+import os
 
-def train_model(model_name, dataset_name, text_column="text"):
-    try:
-        tokenizer = AutoTokenizer.from_pretrained(model_name)
-        model = AutoModelForCausalLM.from_pretrained(model_name)
+# Default model path
+MODEL_DIR = "./custom_model"
 
-        dataset = load_dataset(dataset_name)
+# ---------- Dataset Handling ----------
+def get_dataset(dataset_name, config_name=None, user_file=None):
+    if user_file is not None:
+        # Load user-uploaded text dataset
+        with open(user_file.name, "r", encoding="utf-8") as f:
+            text_data = f.read().splitlines()
+        return Dataset.from_dict({"text": text_data})
 
-        def tokenize(batch):
-            return tokenizer(batch[text_column], padding="max_length", truncation=True)
+    if config_name:
+        return load_dataset(dataset_name, config_name)
+    else:
+        return load_dataset(dataset_name)
 
-        tokenized = dataset.map(tokenize, batched=True)
+# ---------- Training ----------
+def train_model(model_name, dataset_name, config_name, user_file, epochs, output_dir):
+    # Load tokenizer & model
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    model = AutoModelForCausalLM.from_pretrained(model_name)
 
-        training_args = TrainingArguments(
-            output_dir="./results",
-            num_train_epochs=1,
-            per_device_train_batch_size=2,
-            save_steps=10,
-            save_total_limit=1
-        )
+    # Load dataset
+    dataset = get_dataset(dataset_name, config_name, user_file)
 
-        trainer = Trainer(
-            model=model,
-            args=training_args,
-            train_dataset=tokenized["train"]
-        )
+    # Tokenize
+    def tokenize_function(examples):
+        return tokenizer(examples["text"], truncation=True, padding="max_length", max_length=128)
 
-        trainer.train()
-        model.save_pretrained("./custom_model")
+    tokenized_dataset = dataset.map(tokenize_function, batched=True)
 
-        return f"✅ Training complete! Model saved at ./custom_model"
-    except Exception as e:
-        return f"❌ Error: {str(e)}"
+    # Data collator
+    data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)
 
-with gr.Blocks() as demo:
-    gr.Markdown("# 🤖 AI Model Builder (Hugging Face Space)")
-    model_name = gr.Textbox(value="gpt2", label="Base model")
-    dataset_name = gr.Textbox(value="wikitext", label="Dataset name (from HF Datasets)")
-    train_button = gr.Button("Train Model")
-    output = gr.Textbox(label="Status")
-
-    train_button.click(
-        fn=train_model,
-        inputs=[model_name, dataset_name],
-        outputs=output
+    # Training args
+    training_args = TrainingArguments(
+        output_dir=output_dir,
+        evaluation_strategy="no",
+        learning_rate=2e-5,
+        per_device_train_batch_size=2,
+        num_train_epochs=int(epochs),
+        weight_decay=0.01,
+        save_total_limit=1,
+        logging_steps=5,
+    )
+
+    # Trainer
+    trainer = Trainer(
+        model=model,
+        args=training_args,
+        train_dataset=tokenized_dataset["train"] if "train" in tokenized_dataset else tokenized_dataset,
+        tokenizer=tokenizer,
+        data_collator=data_collator,
     )
 
+    trainer.train()
+    trainer.save_model(output_dir)
+    tokenizer.save_pretrained(output_dir)
+
+    return f"✅ Training complete! Model saved to {output_dir}"
+
+# ---------- Chat ----------
+def chat_with_model(prompt, history):
+    if not os.path.exists(MODEL_DIR):
+        return "⚠️ No trained model found yet. Train one first!"
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_DIR)
+    model = AutoModelForCausalLM.from_pretrained(MODEL_DIR)
+
+    inputs = tokenizer(prompt, return_tensors="pt")
+    with torch.no_grad():
+        outputs = model.generate(**inputs, max_length=200, pad_token_id=tokenizer.eos_token_id)
+
+    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return response
+
+# ---------- UI ----------
+with gr.Blocks() as demo:
+    gr.Markdown("# 🧠 Custom AI Model Builder")
+    gr.Markdown("Train and chat with your **own model** directly in Hugging Face.")
+
+    with gr.Tab("Train Model"):
+        model_name = gr.Textbox(label="Base Model (e.g. gpt2, distilgpt2, codeparrot-small)", value="distilgpt2")
+        dataset_name = gr.Textbox(label="Dataset Name (HuggingFace hub, e.g. wikitext, imdb)", value="wikitext")
+        config_name = gr.Textbox(label="Config (optional, e.g. wikitext-2-raw-v1)", value="wikitext-2-raw-v1")
+        user_file = gr.File(label="Or Upload Your Own TXT Dataset", file_types=[".txt"], type="file")
+        epochs = gr.Number(label="Epochs", value=1, precision=0)
+        output_dir = gr.Textbox(label="Output Directory", value=MODEL_DIR)
+        train_button = gr.Button("🚀 Start Training")
+        train_output = gr.Textbox(label="Training Logs")
+
+        train_button.click(
+            train_model,
+            inputs=[model_name, dataset_name, config_name, user_file, epochs, output_dir],
+            outputs=train_output,
+        )
+
+    with gr.Tab("Chat with Model"):
+        chatbot = gr.Chatbot()
+        msg = gr.Textbox(label="Message")
+        send = gr.Button("Send")
+
+        def respond(message, chat_history):
+            response = chat_with_model(message, chat_history)
+            chat_history.append((message, response))
+            return "", chat_history
+
+        send.click(respond, [msg, chatbot], [msg, chatbot])
+
 demo.launch()