app.py

# app.py (Part 1 of 2)

import os
import json
import datetime
import gradio as gr
from transformers import (
    AutoTokenizer, AutoModelForSequenceClassification,
    Trainer, TrainingArguments
)
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):
    safe_name = model_name.replace("/", "_")
    return os.path.join("memory", f"memory_{safe_name}.json")

def load_memory(model_name):
    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, 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):
    try:
        log.info(f"Loading dataset: {dataset_name}")
        parts = dataset_name.split(" ")
        
        if len(parts) == 2:
            dataset_repo, dataset_config = parts
            dataset = load_dataset(dataset_repo, dataset_config, split="train[:200]")  # CPU-friendly subset
        else:
            dataset = load_dataset(dataset_name, split="train[:200]")

        log.info("Dataset loaded successfully")

        tokenizer = AutoTokenizer.from_pretrained(model_name)
        model = AutoModelForSequenceClassification.from_pretrained(model_name, num_labels=2)

        def tokenize_fn(examples):
            return tokenizer(examples["text"], truncation=True, padding="max_length", max_length=128)

        dataset = dataset.map(tokenize_fn, batched=True)
        dataset.set_format(type="torch", columns=["input_ids", "attention_mask", "label"])

        training_args = TrainingArguments(
            output_dir="./results",
            eval_strategy="epoch",
            save_strategy="epoch",
            learning_rate=2e-5,
            per_device_train_batch_size=4,
            num_train_epochs=int(epochs),
            logging_dir="./logs",
            logging_steps=10,
        )

        trainer = Trainer(
            model=model,
            args=training_args,
            train_dataset=dataset,
            tokenizer=tokenizer,
        )

        trainer.train()
        return "Training complete ✅"
    except Exception as e:
        log.error(f"Training failed: {e}")
        return f"Error during training: {e}"

    # 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
    )

    # 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]
    )

    trainer.train()

    # Save trained model
    save_dir = os.path.join(output_dir, model_name.replace("/", "_"))
    model.save_pretrained(save_dir)
    tokenizer.save_pretrained(save_dir)

    log_event(f"Training finished: model saved at {save_dir}")
    return f"✅ Training complete. Model saved at {save_dir}", progress

    # app.py (Part 2 of 2) — UI

import gradio as gr

# =========================
# 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()