Update app.py

app.py

@@ -1,289 +1,286 @@
+# app.py (Part 1 of 2)
+
 import os
 import json
-import logging
-from datetime import datetime
+import datetime
 import gradio as gr
-from datasets import list_datasets, load_dataset
-from huggingface_hub import HfApi, HfFolder
 from transformers import (
-    AutoModelForCausalLM,
-    AutoTokenizer,
-    Trainer,
-    TrainingArguments,
-    DataCollatorForLanguageModeling
-)
-
-# ===============================
-# Setup directories & logging
-# ===============================
-BASE_DIR = "storage"
-MEMORY_DIR = os.path.join(BASE_DIR, "memory")
-LOG_FILE = os.path.join(BASE_DIR, "logs.txt")
-
-os.makedirs(MEMORY_DIR, exist_ok=True)
-os.makedirs(BASE_DIR, exist_ok=True)
-
-logging.basicConfig(
-    filename=LOG_FILE,
-    level=logging.INFO,
-    format="%(asctime)s - %(levelname)s - %(message)s"
+    AutoTokenizer, AutoModelForSequenceClassification,
+    Trainer, TrainingArguments
 )
-
-def log_event(event: str):
-    """Log both to file and console"""
-    logging.info(event)
-    print(event)
-
-# ===============================
-# Hugging Face Auto-fetch
-# ===============================
-def fetch_top_models(limit=10):
-    """Fetch top models from Hugging Face Hub"""
-    api = HfApi()
-    models = api.list_models(sort="downloads", limit=limit)
-    return [m.modelId for m in models]
-
-def fetch_top_datasets(limit=10):
-    """Fetch top datasets from Hugging Face Hub"""
-    api = HfApi()
-    datasets = api.list_datasets(sort="downloads", limit=limit)
-    return [d.id for d in datasets]
-
-TOP_MODELS = fetch_top_models()
-TOP_DATASETS = fetch_top_datasets()
-
-# ===============================
-# Memory Management
-# ===============================
+from datasets import load_dataset
+import torch
+
+# =========================
+# Ensure directories exist
+# =========================
+os.makedirs("trained_models", exist_ok=True)
+os.makedirs("logs", exist_ok=True)
+os.makedirs("memory", exist_ok=True)
+
+# =========================
+# Utility: Memory System
+# =========================
 def get_memory_file(model_name):
-    return os.path.join(MEMORY_DIR, f"{model_name.replace('/', '_')}_memory.json")
+    safe_name = model_name.replace("/", "_")
+    return os.path.join("memory", f"memory_{safe_name}.json")
 
 def load_memory(model_name):
-    f = get_memory_file(model_name)
-    if os.path.exists(f):
-        with open(f, "r") as file:
-            return json.load(file)
+    file = get_memory_file(model_name)
+    if os.path.exists(file):
+        with open(file, "r") as f:
+            return json.load(f)
     return []
 
-def save_memory(model_name, messages):
-    f = get_memory_file(model_name)
-    with open(f, "w") as file:
-        json.dump(messages, file, indent=2)
-
-# ===============================
-# Chat Functionality
-# ===============================
-def chat_with_model(user_input, model_choice):
-    if not model_choice:
-        return "❌ Please select a model.", ""
-
-    log_event(f"User chatting with {model_choice}: {user_input}")
-    tokenizer = AutoTokenizer.from_pretrained(model_choice)
-    model = AutoModelForCausalLM.from_pretrained(model_choice)
-
-    inputs = tokenizer(user_input, return_tensors="pt")
-    outputs = model.generate(**inputs, max_length=200)
-
-    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-
-    memory = load_memory(model_choice)
-    memory.append({"user": user_input, "bot": response})
-    save_memory(model_choice, memory)
-
-    return response, json.dumps(memory, indent=2)
-
-# ===============================
-# Training
-# ===============================
-def train_model(model_name, dataset_name, epochs, output_dir):
-    try:
-        log_event(f"Starting training: model={model_name}, dataset={dataset_name}, epochs={epochs}")
-        dataset = load_dataset(dataset_name, split="train")
-
-        tokenizer = AutoTokenizer.from_pretrained(model_name)
-        model = AutoModelForCausalLM.from_pretrained(model_name)
-
-        def tokenize_function(examples):
-            return tokenizer(examples["text"], truncation=True, padding="max_length", max_length=128)
-
-        tokenized_dataset = dataset.map(tokenize_function, batched=True)
-
-        data_collator = DataCollatorForLanguageModeling(
-            tokenizer=tokenizer,
-            mlm=False
-        )
-
-        training_args = TrainingArguments(
-            output_dir=output_dir,
-            overwrite_output_dir=True,
-            num_train_epochs=int(epochs),
-            per_device_train_batch_size=2,
-            save_steps=500,
-            save_total_limit=2,
-            logging_dir="./logs",
-            logging_steps=50
-        )
-
-        trainer = Trainer(
-            model=model,
-            args=training_args,
-            train_dataset=tokenized_dataset,
-            data_collator=data_collator
-        )
-
-        trainer.train()
-        model.save_pretrained(output_dir)
-        tokenizer.save_pretrained(output_dir)
-
-        log_event(f"✅ Training completed. Model saved to {output_dir}")
-        return f"✅ Training completed. Model saved to {output_dir}"
-    except Exception as e:
-        log_event(f"❌ Training failed: {e}")
-        return f"❌ Error during training: {str(e)}"
-
-# ===============================
-# Gradio UI – Training Tab
-# ===============================
-with gr.Blocks() as training_tab:
-    gr.Markdown("## 📚 Train a Custom Model")
-    with gr.Row():
-        model_dropdown = gr.Dropdown(choices=TOP_MODELS, label="Choose Model", interactive=True)
-        dataset_dropdown = gr.Dropdown(choices=TOP_DATASETS, label="Choose Dataset", interactive=True)
-    with gr.Row():
-        model_text = gr.Textbox(label="Or enter custom model ID", placeholder="e.g. gpt2")
-        dataset_text = gr.Textbox(label="Or enter custom dataset ID", placeholder="e.g. wikitext")
-    epochs = gr.Number(value=1, label="Epochs")
-    output_dir = gr.Textbox(value="./trained_model", label="Output Directory")
-    train_btn = gr.Button("🚀 Start Training")
-    train_output = gr.Textbox(label="Training Status")
-
-    def handle_train(model_d, model_t, dataset_d, dataset_t, epochs, output_dir):
-        model = model_t if model_t else model_d
-        dataset = dataset_t if dataset_t else dataset_d
-        return train_model(model, dataset, epochs, output_dir)
-
-    train_btn.click(
-        fn=handle_train,
-        inputs=[model_dropdown, model_text, dataset_dropdown, dataset_text, epochs, output_dir],
-        outputs=train_output
+def save_memory(model_name, conversation):
+    file = get_memory_file(model_name)
+    memory = load_memory(model_name)
+    memory.append(conversation)
+    with open(file, "w") as f:
+        json.dump(memory, f, indent=2)
+
+# =========================
+# Utility: Logging
+# =========================
+def log_event(event):
+    log_file = os.path.join("logs", "events.log")
+    with open(log_file, "a") as f:
+        f.write(f"[{datetime.datetime.now()}] {event}\n")
+
+# =========================
+# Training Pipeline
+# =========================
+def train_model(model_name, dataset_name, epochs, output_dir="trained_models"):
+    log_event(f"Training started: model={model_name}, dataset={dataset_name}, epochs={epochs}")
+
+    # Load tokenizer + dataset
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    dataset = load_dataset(dataset_name, split="train[:200]")  # smaller subset for CPU
+
+    def tokenize_fn(batch):
+        return tokenizer(batch["text"], truncation=True, padding="max_length", max_length=128)
+
+    dataset = dataset.map(tokenize_fn, batched=True)
+    dataset = dataset.rename_column("label", "labels")
+    dataset.set_format("torch", columns=["input_ids", "attention_mask", "labels"])
+
+    # Load model
+    model = AutoModelForSequenceClassification.from_pretrained(model_name, num_labels=2)
+
+    # Training arguments
+    training_args = TrainingArguments(
+        output_dir=os.path.join(output_dir, model_name.replace("/", "_")),
+        overwrite_output_dir=True,
+        evaluation_strategy="epoch",
+        save_strategy="epoch",
+        num_train_epochs=epochs,
+        per_device_train_batch_size=8,
+        logging_dir="./logs",
+        logging_steps=10,
+        report_to="none",  # prevent wandb errors
+        no_cuda=True  # force CPU
     )
 
-    # ===============================
-# Gradio UI – Chat Tab
-# ===============================
-with gr.Blocks() as chat_tab:
-    gr.Markdown("## 💬 Chat with Model")
-    with gr.Row():
-        chat_model_dropdown = gr.Dropdown(choices=TOP_MODELS, label="Choose Model", interactive=True)
-        chat_model_text = gr.Textbox(label="Or enter custom model ID", placeholder="e.g. gpt2")
-    with gr.Row():
-        chat_input = gr.Textbox(label="Your Message")
-        send_btn = gr.Button("Send")
-    chat_output = gr.Textbox(label="Model Response")
-    memory_display = gr.Textbox(label="Conversation Memory", interactive=False)
-
-    def handle_chat(user_input, model_d, model_t):
-        model = model_t if model_t else model_d
-        return chat_with_model(user_input, model)
-
-    send_btn.click(
-        fn=handle_chat,
-        inputs=[chat_input, chat_model_dropdown, chat_model_text],
-        outputs=[chat_output, memory_display]
+    # Progress tracking
+    progress = {"epoch": 0, "loss": []}
+
+    def compute_metrics(eval_pred):
+        logits, labels = eval_pred
+        preds = logits.argmax(-1)
+        acc = (preds == labels).astype(float).mean().item()
+        return {"accuracy": acc}
+
+    def log_callback(trainer, state, control, **kwargs):
+        if state.is_local_process_zero and state.log_history:
+            last_log = state.log_history[-1]
+            if "loss" in last_log:
+                progress["epoch"] = state.epoch
+                progress["loss"].append(last_log["loss"])
+                log_event(f"Epoch {state.epoch} - Loss: {last_log['loss']}")
+
+    # Trainer
+    trainer = Trainer(
+        model=model,
+        args=training_args,
+        train_dataset=dataset,
+        tokenizer=tokenizer,
+        compute_metrics=compute_metrics,
+        callbacks=[log_callback]
     )
 
-# ===============================
-# Gradio UI – Memory Tab
-# ===============================
-with gr.Blocks() as memory_tab:
-    gr.Markdown("## 🧠 Manage Memory")
-
-    with gr.Row():
-        memory_model_dropdown = gr.Dropdown(choices=TOP_MODELS, label="Select Model")
-        memory_model_text = gr.Textbox(label="Or enter custom model ID")
-
-    memory_output = gr.Textbox(label="Stored Memory", interactive=False)
-    load_btn = gr.Button("📂 Load Memory")
-    clear_btn = gr.Button("🗑️ Clear Memory")
-
-    def handle_load(model_d, model_t):
-        model = model_t if model_t else model_d
-        memory = load_memory(model)
-        return json.dumps(memory, indent=2)
-
-    def handle_clear(model_d, model_t):
-        model = model_t if model_t else model_d
-        f = get_memory_file(model)
-        if os.path.exists(f):
-            os.remove(f)
-            log_event(f"Cleared memory for {model}")
-            return "✅ Memory cleared."
-        return "⚠️ No memory found."
-
-    load_btn.click(
-        fn=handle_load,
-        inputs=[memory_model_dropdown, memory_model_text],
-        outputs=memory_output
-    )
-
-    clear_btn.click(
-        fn=handle_clear,
-        inputs=[memory_model_dropdown, memory_model_text],
-        outputs=memory_output
-    )
+    trainer.train()
 
-# ===============================
-# Gradio UI – Logs Tab
-# ===============================
-with gr.Blocks() as logs_tab:
-    gr.Markdown("## 📜 Application Logs")
-    log_display = gr.Textbox(value=open(LOG_FILE).read() if os.path.exists(LOG_FILE) else "No logs yet.", lines=20)
+    # Save trained model
+    save_dir = os.path.join(output_dir, model_name.replace("/", "_"))
+    model.save_pretrained(save_dir)
+    tokenizer.save_pretrained(save_dir)
 
-    refresh_btn = gr.Button("🔄 Refresh Logs")
+    log_event(f"Training finished: model saved at {save_dir}")
+    return f"✅ Training complete. Model saved at {save_dir}", progress
 
-    def refresh_logs():
-        return open(LOG_FILE).read() if os.path.exists(LOG_FILE) else "No logs yet."
+    # app.py (Part 2 of 2) — UI
 
-    refresh_btn.click(
-        fn=refresh_logs,
-        outputs=log_display
-    )
+import gradio as gr
 
-# ===============================
-# Gradio UI – Help Tab
-# ===============================
-with gr.Blocks() as help_tab:
-    gr.Markdown("## 📖 Help & User Manual")
-
-    gr.Markdown("""
-### 🔹 Beginner Guide
-1. Go to **Train a Model** tab → pick a model & dataset or enter custom IDs.  
-2. Choose number of epochs & output directory → click **Start Training**.  
-3. Once training completes, the model is saved and can be used later.  
-4. Go to **Chat with Model** tab → type your message or test the trained model.  
-5. Conversation is auto-saved per model (see **Memory** tab).  
-6. Use **Logs** tab for detailed runtime events.  
-
-### 🔹 Technical Details
-- Models/Datasets pulled live from Hugging Face Hub (top 10 auto-fetched).  
-- Memory stored in `/storage/memory/` as JSON files (per model).  
-- Logs stored in `/storage/logs.txt`.  
-- Training uses 🤗 Transformers `Trainer` API with causal LM objective.  
-- Safe checks auto-create missing directories & files.  
-
-### 🔹 Tips
-- For large datasets, train on GPU (CPU will be very slow).  
-- Memory files can be manually edited in `/storage/memory/`.  
-- You can load any public Hugging Face dataset/model by entering its ID.  
-    """)
-
-# ===============================
-# Final Tabbed Interface
-# ===============================
-demo = gr.TabbedInterface(
-    [training_tab, chat_tab, memory_tab, logs_tab, help_tab],
-    ["Train a Model", "Chat", "Memory", "Logs", "Help & Manual"]
-)
+# =========================
+# Hugging Face Top 10 (demo defaults, can expand to auto-fetch later)
+# =========================
+TOP_MODELS = [
+    "distilbert-base-uncased", "bert-base-uncased", "roberta-base",
+    "google/electra-base-discriminator", "albert-base-v2",
+    "facebook/bart-base", "gpt2", "t5-small",
+    "microsoft/deberta-base", "xlnet-base-cased"
+]
+
+TOP_DATASETS = [
+    "imdb", "ag_news", "yelp_polarity",
+    "dbpedia_14", "amazon_polarity",
+    "tweet_eval", "glue", "sst2",
+    "cnn_dailymail", "emotion"
+]
+
+# =========================
+# Inference (Test chat)
+# =========================
+def chat_with_model(model_name, user_input):
+    model_dir = os.path.join("trained_models", model_name.replace("/", "_"))
+    if not os.path.exists(model_dir):
+        return "❌ Model not trained yet. Train it first."
+
+    tokenizer = AutoTokenizer.from_pretrained(model_dir)
+    model = AutoModelForSequenceClassification.from_pretrained(model_dir)
+
+    inputs = tokenizer(user_input, return_tensors="pt", truncation=True, padding=True)
+    outputs = model(**inputs)
+    prediction = torch.argmax(outputs.logits, dim=-1).item()
+
+    # Save memory
+    conversation = {"input": user_input, "prediction": prediction}
+    save_memory(model_name, conversation)
+
+    return f"🔮 Prediction: {prediction}"
+
+# =========================
+# View Memory
+# =========================
+def view_memory(model_name):
+    memory = load_memory(model_name)
+    if not memory:
+        return "📭 No memory yet for this model."
+    return json.dumps(memory, indent=2)
+
+# =========================
+# View Logs
+# =========================
+def view_logs():
+    log_file = os.path.join("logs", "events.log")
+    if not os.path.exists(log_file):
+        return "📭 No logs yet."
+    with open(log_file, "r") as f:
+        return f.read()
+
+# =========================
+# User Guide / Manual
+# =========================
+USER_GUIDE = """
+# 📘 AI Model Builder Guide
+
+Welcome to your **all-in-one AI Model Builder**.  
+This app allows you to **train, fine-tune, test, and manage AI models** directly in a Hugging Face Space.  
+
+---
+
+## 🔹 Step 1: Training a Model
+1. Go to the **Training Tab**.
+2. Select a **model** from the Top-10 list or type your own Hugging Face model ID.
+3. Select a **dataset** from the Top-10 list or type your own Hugging Face dataset ID.
+4. Choose the number of **epochs** (training cycles).
+5. Click **Start Training**.
+6. Training progress will appear, and the model will be saved under `trained_models/`.
+
+---
+
+## 🔹 Step 2: Testing Your Model
+1. Switch to the **Testing Tab**.
+2. Type any input in the chat box.
+3. The app will return a **prediction**.
+4. Every conversation is saved in **per-model memory**.
+
+---
+
+## 🔹 Step 3: Viewing Memory
+- Go to the **Memory Tab**.
+- See past chats and predictions for each model.
+
+---
+
+## 🔹 Step 4: Viewing Logs
+- All activity is logged.
+- Open the **Logs Tab** to view training sessions, progress, and errors.
+
+---
+
+## 🔹 Technical Notes
+- Training runs on **CPU** (slower but free).
+- Uses Hugging Face **Transformers + Datasets**.
+- Stores:
+  - Models → `trained_models/`
+  - Logs → `logs/events.log`
+  - Memory → `memory/memory_{model}.json`
+"""
+
+# =========================
+# Build Gradio UI
+# =========================
+with gr.Blocks() as demo:
+    gr.Markdown("# 🧠 AI Model Builder\nTrain, Fine-tune, Test, and Manage Your Own AI Models")
+
+    with gr.Tab("🛠️ Training"):
+        with gr.Row():
+            model_dropdown = gr.Dropdown(choices=TOP_MODELS, label="Select Model", interactive=True)
+            model_textbox = gr.Textbox(label="Or enter custom model ID")
+        with gr.Row():
+            dataset_dropdown = gr.Dropdown(choices=TOP_DATASETS, label="Select Dataset", interactive=True)
+            dataset_textbox = gr.Textbox(label="Or enter custom dataset ID")
+        epochs = gr.Slider(1, 5, value=1, step=1, label="Epochs (Training Cycles)")
+        train_button = gr.Button("🚀 Start Training")
+        train_output = gr.Textbox(label="Training Status")
+        progress_output = gr.JSON(label="Progress Details")
+
+        def run_training(model_dropdown, model_textbox, dataset_dropdown, dataset_textbox, epochs):
+            model_name = model_textbox if model_textbox else model_dropdown
+            dataset_name = dataset_textbox if dataset_textbox else dataset_dropdown
+            return train_model(model_name, dataset_name, epochs)
+
+        train_button.click(
+            run_training,
+            inputs=[model_dropdown, model_textbox, dataset_dropdown, dataset_textbox, epochs],
+            outputs=[train_output, progress_output]
+        )
 
+    with gr.Tab("💬 Testing"):
+        test_model_name = gr.Textbox(label="Enter Model ID (must be trained first)")
+        test_input = gr.Textbox(label="Your Message")
+        test_button = gr.Button("💡 Predict")
+        test_output = gr.Textbox(label="Model Response")
+        test_button.click(chat_with_model, inputs=[test_model_name, test_input], outputs=test_output)
+
+    with gr.Tab("🧾 Memory"):
+        mem_model_name = gr.Textbox(label="Enter Model ID to View Memory")
+        mem_button = gr.Button("📂 Load Memory")
+        mem_output = gr.Textbox(label="Conversation Memory", lines=15)
+        mem_button.click(view_memory, inputs=mem_model_name, outputs=mem_output)
+
+    with gr.Tab("📜 Logs"):
+        log_button = gr.Button("📖 Show Logs")
+        log_output = gr.Textbox(label="Logs", lines=20)
+        log_button.click(view_logs, outputs=log_output)
+
+    with gr.Tab("📘 Guide"):
+        gr.Markdown(USER_GUIDE)
+
+# =========================
+# Launch
+# =========================
 if __name__ == "__main__":
-    demo.launch(server_name="0.0.0.0", server_port=7860)
+    demo.launch()