Update app.py (#11)

- Update app.py (29f771d050edf3bcff1d9c69395fcf7877857223)

app.py

@@ -1,8 +1,10 @@
 import os
 import json
+import logging
+from datetime import datetime
 import gradio as gr
-from huggingface_hub import HfApi, snapshot_download
-from datasets import load_dataset
+from datasets import list_datasets, load_dataset
+from huggingface_hub import HfApi, HfFolder
 from transformers import (
     AutoModelForCausalLM,
     AutoTokenizer,
@@ -10,184 +12,278 @@ from transformers import (
     TrainingArguments,
     DataCollatorForLanguageModeling
 )
-import torch
 
-# ========== GLOBALS ==========
-API = HfApi()
-MEMORY_DIR = "memory"
-MODEL_DIR = "saved_models"
+# ===============================
+# 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(MODEL_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"
+)
+
+def log_event(event: str):
+    """Log both to file and console"""
+    logging.info(event)
+    print(event)
 
-# ========== MEMORY HANDLING ==========
+# ===============================
+# 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
+# ===============================
 def get_memory_file(model_name):
-    return os.path.join(MEMORY_DIR, f"{model_name}_memory.json")
+    return os.path.join(MEMORY_DIR, f"{model_name.replace('/', '_')}_memory.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)
+    f = get_memory_file(model_name)
+    if os.path.exists(f):
+        with open(f, "r") as file:
+            return json.load(file)
     return []
 
-def save_memory(model_name, chat_log):
-    file = get_memory_file(model_name)
-    with open(file, "w") as f:
-        json.dump(chat_log, f, indent=2)
+def save_memory(model_name, messages):
+    f = get_memory_file(model_name)
+    with open(f, "w") as file:
+        json.dump(messages, file, indent=2)
 
-# ========== HUGGING FACE HUB HELPERS ==========
-def get_top_models(limit=10):
-    models = API.list_models(sort="downloads", direction=-1, limit=limit)
-    return [m.modelId for m in models]
+# ===============================
+# Chat Functionality
+# ===============================
+def chat_with_model(user_input, model_choice):
+    if not model_choice:
+        return "❌ Please select a model.", ""
 
-def get_top_datasets(limit=10):
-    datasets = API.list_datasets(sort="downloads", direction=-1, limit=limit)
-    return [d.id for d in datasets]
+    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)
 
-# ========== TRAINING ==========
-def train_model(model_name, dataset_name, output_name, epochs=1):
+    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:
-        dataset = load_dataset(dataset_name)
+        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_fn(examples):
-            return tokenizer(examples["text"], truncation=True, padding="max_length")
+        def tokenize_function(examples):
+            return tokenizer(examples["text"], truncation=True, padding="max_length", max_length=128)
 
-        tokenized = dataset.map(tokenize_fn, batched=True, remove_columns=dataset["train"].column_names)
-        collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)
+        tokenized_dataset = dataset.map(tokenize_function, batched=True)
+
+        data_collator = DataCollatorForLanguageModeling(
+            tokenizer=tokenizer,
+            mlm=False
+        )
 
         training_args = TrainingArguments(
-            output_dir=os.path.join(MODEL_DIR, output_name),
+            output_dir=output_dir,
             overwrite_output_dir=True,
-            num_train_epochs=epochs,
+            num_train_epochs=int(epochs),
             per_device_train_batch_size=2,
-            save_strategy="epoch",
+            save_steps=500,
+            save_total_limit=2,
             logging_dir="./logs",
-            logging_steps=10,
-            push_to_hub=False
+            logging_steps=50
         )
 
         trainer = Trainer(
             model=model,
             args=training_args,
-            train_dataset=tokenized["train"],
-            tokenizer=tokenizer,
-            data_collator=collator
+            train_dataset=tokenized_dataset,
+            data_collator=data_collator
         )
 
         trainer.train()
-        trainer.save_model(os.path.join(MODEL_DIR, output_name))
-        return f"✅ Training finished. Model saved as {output_name}"
+        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:
-        return f"❌ Training error: {str(e)}"
+        log_event(f"❌ Training failed: {e}")
+        return f"❌ Error during training: {str(e)}"
 
-# ========== CHAT ==========
-def chat_with_model(model_name, user_input):
-    try:
-        path = os.path.join(MODEL_DIR, model_name)
-        if os.path.exists(path):
-            model = AutoModelForCausalLM.from_pretrained(path)
-            tokenizer = AutoTokenizer.from_pretrained(path)
-        else:
-            model = AutoModelForCausalLM.from_pretrained(model_name)
-            tokenizer = AutoTokenizer.from_pretrained(model_name)
+# ===============================
+# 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")
 
-        memory = load_memory(model_name)
-        memory_text = " ".join([f"User: {m['user']} AI: {m['ai']}" for m in memory])
+    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)
 
-        inputs = tokenizer(memory_text + " User: " + user_input + " AI:", return_tensors="pt")
-        outputs = model.generate(**inputs, max_length=300, pad_token_id=tokenizer.eos_token_id)
-        response = tokenizer.decode(outputs[0], skip_special_tokens=True).split("AI:")[-1].strip()
+    train_btn.click(
+        fn=handle_train,
+        inputs=[model_dropdown, model_text, dataset_dropdown, dataset_text, epochs, output_dir],
+        outputs=train_output
+    )
 
-        memory.append({"user": user_input, "ai": response})
-        save_memory(model_name, memory)
+    # ===============================
+# 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)
 
-        return response
-    except Exception as e:
-        return f"❌ Chat error: {str(e)}"
+    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)
 
-# ========== INTERFACE TABS ==========
+    send_btn.click(
+        fn=handle_chat,
+        inputs=[chat_input, chat_model_dropdown, chat_model_text],
+        outputs=[chat_output, memory_display]
+    )
 
-# Home Tab
-home_tab = gr.Markdown("# 🤖 My AI Model Builder\nWelcome! Use this tool to search, train, and chat with AI models. All memory & models are stored locally inside this Space.")
+# ===============================
+# Gradio UI – Memory Tab
+# ===============================
+with gr.Blocks() as memory_tab:
+    gr.Markdown("## 🧠 Manage Memory")
 
-# Models Tab
-with gr.Blocks() as models_tab:
-    gr.Markdown("## 🔍 Search Models")
     with gr.Row():
-        model_input = gr.Textbox(label="Search or enter model name")
-        model_output = gr.Textbox(label="Result")
-    search_button = gr.Button("Search Model")
-
-    def search_model(name):
-        try:
-            info = API.model_info(name)
-            return f"✅ Found: {info.modelId}\nDownloads: {info.downloads}\nTags: {info.tags}"
-        except Exception as e:
-            return f"❌ {str(e)}"
-    search_button.click(search_model, model_input, model_output)
-
-# Datasets Tab
-with gr.Blocks() as datasets_tab:
-    gr.Markdown("## 📊 Search Datasets")
-    with gr.Row():
-        dataset_input = gr.Textbox(label="Search or enter dataset name")
-        dataset_output = gr.Textbox(label="Result")
-    dataset_button = gr.Button("Search Dataset")
-
-    def search_dataset(name):
-        try:
-            info = API.dataset_info(name)
-            return f"✅ Found: {info.id}\nDownloads: {info.downloads}\nTags: {info.tags}"
-        except Exception as e:
-            return f"❌ {str(e)}"
-    dataset_button.click(search_dataset, dataset_input, dataset_output)
-
-# Training Tab
-with gr.Blocks() as training_tab:
-    gr.Markdown("## 🏋️ Train / Fine-tune a Model")
-    model_choice = gr.Dropdown(choices=get_top_models(), label="Pick Base Model", interactive=True)
-    dataset_choice = gr.Dropdown(choices=get_top_datasets(), label="Pick Dataset", interactive=True)
-    output_name = gr.Textbox(label="New Model Name")
-    epochs = gr.Slider(1, 5, step=1, label="Epochs", value=1)
-    train_button = gr.Button("🚀 Train Model")
-    train_output = gr.Textbox(label="Training Status")
-    train_button.click(train_model, [model_choice, dataset_choice, output_name, epochs], train_output)
+        memory_model_dropdown = gr.Dropdown(choices=TOP_MODELS, label="Select Model")
+        memory_model_text = gr.Textbox(label="Or enter custom model ID")
 
-# Chat/Test Tab
-with gr.Blocks() as chat_tab:
-    gr.Markdown("## 💬 Chat with Your Model")
-    chat_model = gr.Textbox(label="Enter Model Name")
-    user_input = gr.Textbox(label="Your Message")
-    chat_output = gr.Textbox(label="AI Response")
-    chat_button = gr.Button("Send")
-    chat_button.click(chat_with_model, [chat_model, user_input], chat_output)
-
-# Memory Tab
-with gr.Blocks() as memory_tab:
-    gr.Markdown("## 🧠 Model Memory")
-    memory_model = gr.Textbox(label="Model Name")
-    memory_display = gr.Textbox(label="Memory Log")
-    def show_memory(name): return json.dumps(load_memory(name), indent=2)
-    memory_button = gr.Button("Load Memory")
-    memory_button.click(show_memory, memory_model, memory_display)
-
-# Guide Tab
-guide_tab = gr.Markdown("""
-# 📖 Mini Guide
-1. Use **Models** to explore Hugging Face models.
-2. Use **Datasets** to find training data.
-3. Use **Training** to fine-tune.
-4. Use **Chat** to test models.
-5. All models & memory are saved in this Space.
-""")
-
-# Launch Interface
+    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
+    )
+
+# ===============================
+# 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)
+
+    refresh_btn = gr.Button("🔄 Refresh Logs")
+
+    def refresh_logs():
+        return open(LOG_FILE).read() if os.path.exists(LOG_FILE) else "No logs yet."
+
+    refresh_btn.click(
+        fn=refresh_logs,
+        outputs=log_display
+    )
+
+# ===============================
+# 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(
-    [home_tab, models_tab, datasets_tab, training_tab, chat_tab, memory_tab, guide_tab],
-    ["Home", "Models", "Datasets", "Training", "Chat/Test", "Memory", "Guide"]
+    [training_tab, chat_tab, memory_tab, logs_tab, help_tab],
+    ["Train a Model", "Chat", "Memory", "Logs", "Help & Manual"]
 )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch(server_name="0.0.0.0", server_port=7860)