Update app.py (#6)

- Update app.py (735d33ea0ad6303d462d059ba812cbef76691f83)

app.py

@@ -1,124 +1,156 @@
+import os
+import json
 import gradio as gr
+from datetime import datetime
+from transformers import AutoTokenizer, AutoModelForCausalLM, Trainer, TrainingArguments
 from datasets import load_dataset
-from transformers import (
-    AutoTokenizer,
-    AutoModelForCausalLM,
-    Trainer,
-    TrainingArguments,
-    DataCollatorForLanguageModeling,
-)
-import torch
-
-
-# Map specialization → dataset + base model
-SPECIALIZATIONS = {
-    "Coding Assistant": {
-        "dataset": "codeparrot/github-code",
-        "model": "EleutherAI/gpt-neo-125M",
-    },
-    "Cybersecurity Helper": {
-        "dataset": "wikitext",
-        "model": "distilgpt2",  # placeholder dataset, replace with cybersecurity text later
-    },
-    "App/Web Developer": {
-        "dataset": "wikitext",
-        "model": "gpt2",
-    },
-    "General Problem Solver": {
-        "dataset": "wikitext",
-        "model": "gpt2",
-    },
-}
-
-
-def train_model(specialization, epochs, lr):
-    try:
-        spec = SPECIALIZATIONS.get(specialization, SPECIALIZATIONS["General Problem Solver"])
-
-        dataset_name = spec["dataset"]
-        model_name = spec["model"]
 
-        # Load dataset
+# ========= MEMORY MANAGEMENT =========
+MEMORY_DIR = "memories"
+MODEL_DIR = "models"
+
+os.makedirs(MEMORY_DIR, exist_ok=True)
+os.makedirs(MODEL_DIR, exist_ok=True)
+
+def get_memory_file(model_name):
+    safe_name = model_name.replace("/", "_")
+    return os.path.join(MEMORY_DIR, f"{safe_name}_memory.json")
+
+def load_memory(model_name):
+    filepath = get_memory_file(model_name)
+    if os.path.exists(filepath):
+        with open(filepath, "r") as f:
+            return json.load(f)
+    return []
+
+def save_memory(model_name, memory_data):
+    filepath = get_memory_file(model_name)
+    with open(filepath, "w") as f:
+        json.dump(memory_data, f, indent=2)
+
+def append_memory(model_name, role, content):
+    memory = load_memory(model_name)
+    memory.append({
+        "timestamp": datetime.now().isoformat(),
+        "role": role,
+        "content": content
+    })
+    save_memory(model_name, memory)
+
+def clear_memory(model_name):
+    filepath = get_memory_file(model_name)
+    if os.path.exists(filepath):
+        os.remove(filepath)
+    return f"Memory cleared for {model_name}."
+
+def download_memory(model_name):
+    filepath = get_memory_file(model_name)
+    if os.path.exists(filepath):
+        return filepath
+    return None
+
+def upload_memory(model_name, file_obj):
+    if file_obj is None:
+        return "No file uploaded."
+    new_data = json.load(open(file_obj.name))
+    save_memory(model_name, new_data)
+    return f"Memory replaced for {model_name}."
+
+def merge_memory(model_name, file_obj):
+    if file_obj is None:
+        return "No file uploaded."
+    current = load_memory(model_name)
+    new_data = json.load(open(file_obj.name))
+    merged = current + new_data
+    save_memory(model_name, merged)
+    return f"Memory merged for {model_name}."
+
+# ========= MODEL MANAGEMENT =========
+def train_model(model_name, dataset_name, epochs, output_dir):
+    try:
         dataset = load_dataset(dataset_name)
-
-        # Load tokenizer & model
         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_datasets = dataset.map(tokenize_function, batched=True, remove_columns=["text"])
-
-        # Data collator
-        data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)
+        def tokenize(batch):
+            return tokenizer(batch["text"], truncation=True, padding="max_length", max_length=128)
 
-        # Training args
+        dataset = dataset.map(tokenize, batched=True)
         training_args = TrainingArguments(
-            output_dir="./results",
-            eval_strategy="epoch",
-            learning_rate=lr,
+            output_dir=output_dir,
+            overwrite_output_dir=True,
             per_device_train_batch_size=2,
-            per_device_eval_batch_size=2,
-            num_train_epochs=epochs,
-            weight_decay=0.01,
-            save_strategy="no",
-            logging_dir="./logs",
-            logging_steps=10,
-        )
-
-        trainer = Trainer(
-            model=model,
-            args=training_args,
-            train_dataset=tokenized_datasets["train"],
-            eval_dataset=tokenized_datasets["validation"],
-            tokenizer=tokenizer,
-            data_collator=data_collator,
+            num_train_epochs=int(epochs),
+            save_strategy="epoch",
+            logging_dir=f"{output_dir}/logs"
         )
 
+        trainer = Trainer(model=model, args=training_args, train_dataset=dataset["train"])
         trainer.train()
-
-        return f"✅ Training complete for {specialization} model ({model_name}) with {epochs} epochs, lr={lr}"
+        model.save_pretrained(output_dir)
+        tokenizer.save_pretrained(output_dir)
+        return f"Training complete. Model saved to {output_dir}"
     except Exception as e:
-        return f"❌ Error: {str(e)}"
+        return f"Error: {str(e)}"
 
-
-# Inference / Chat Function
-def chat_fn(prompt, specialization):
+def chat_with_model(model_name, prompt):
     try:
-        spec = SPECIALIZATIONS.get(specialization, SPECIALIZATIONS["General Problem Solver"])
-        model_name = spec["model"]
-        tokenizer = AutoTokenizer.from_pretrained(model_name)
-        model = AutoModelForCausalLM.from_pretrained(model_name)
+        model_path = os.path.join(MODEL_DIR, model_name.replace("/", "_"))
+        if os.path.exists(model_path):
+            model = AutoModelForCausalLM.from_pretrained(model_path)
+            tokenizer = AutoTokenizer.from_pretrained(model_path)
+        else:
+            model = AutoModelForCausalLM.from_pretrained(model_name)
+            tokenizer = AutoTokenizer.from_pretrained(model_name)
 
         inputs = tokenizer(prompt, return_tensors="pt")
-        outputs = model.generate(**inputs, max_length=200)
-        return tokenizer.decode(outputs[0], skip_special_tokens=True)
-    except Exception as e:
-        return f"❌ Chat error: {str(e)}"
+        outputs = model.generate(**inputs, max_length=256, do_sample=True, temperature=0.7)
+        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
 
+        append_memory(model_name, "user", prompt)
+        append_memory(model_name, "assistant", response)
 
-# Gradio UI
-with gr.Blocks() as demo:
-    gr.Markdown("# 🚀 Custom AI Model Builder & Assistant")
+        return response
+    except Exception as e:
+        return f"Error: {str(e)}"
 
-    with gr.Tab("1️⃣ Train Custom Model"):
-        specialization = gr.Radio(
-            list(SPECIALIZATIONS.keys()),
-            label="What do you want your AI to specialize in?",
-            value="General Problem Solver",
-        )
-        epochs = gr.Slider(1, 10, value=1, step=1, label="Training Epochs")
-        lr = gr.Slider(1e-6, 5e-4, value=5e-5, step=1e-6, label="Learning Rate")
-        train_button = gr.Button("🚀 Start Training")
-        output_log = gr.Textbox(label="Training Log")
-        train_button.click(train_model, inputs=[specialization, epochs, lr], outputs=output_log)
-
-    with gr.Tab("2️⃣ Chat with Your Model"):
-        chat_specialization = gr.Dropdown(list(SPECIALIZATIONS.keys()), value="General Problem Solver", label="Model Type")
-        prompt = gr.Textbox(label="Ask me anything", placeholder="Type your question here...")
-        chat_button = gr.Button("💬 Generate Response")
+# ========= INTERFACE =========
+with gr.Blocks() as demo:
+    gr.Markdown("# 🤖 My AI Model Builder\nTrain, fine-tune, test, and manage AI models with memory.")
+
+    with gr.Tab("Train Model"):
+        model_name = gr.Textbox(label="Base Model (Hugging Face Hub ID)", value="gpt2")
+        dataset_name = gr.Textbox(label="Dataset Name (Hugging Face Dataset ID)", value="wikitext")
+        epochs = gr.Number(label="Epochs", value=1, precision=0)
+        output_dir = gr.Textbox(label="Output Directory", value="models/custom_model")
+        train_btn = gr.Button("Train Model")
+        train_output = gr.Textbox(label="Training Status")
+        train_btn.click(train_model, inputs=[model_name, dataset_name, epochs, output_dir], outputs=train_output)
+
+    with gr.Tab("Test Models / Chat"):
+        chat_model = gr.Textbox(label="Model Name", value="gpt2")
+        user_prompt = gr.Textbox(label="Enter Prompt")
+        chat_btn = gr.Button("Chat")
         chat_output = gr.Textbox(label="Response")
-        chat_button.click(chat_fn, inputs=[prompt, chat_specialization], outputs=chat_output)
+        chat_btn.click(chat_with_model, inputs=[chat_model, user_prompt], outputs=chat_output)
+
+    with gr.Tab("Memory Management"):
+        mem_model = gr.Textbox(label="Model Name", value="gpt2")
+        view_btn = gr.Button("View Memory")
+        memory_output = gr.JSON(label="Memory Log")
+        view_btn.click(load_memory, inputs=[mem_model], outputs=memory_output)
+
+        with gr.Row():
+            dl_btn = gr.Button("Download Memory")
+            up_btn = gr.File(label="Upload Memory JSON")
+            merge_btn = gr.File(label="Merge Memory JSON")
+
+        dl_file = gr.File()
+        dl_btn.click(download_memory, inputs=[mem_model], outputs=dl_file)
+        up_btn.upload(upload_memory, inputs=[mem_model, up_btn], outputs=memory_output)
+        merge_btn.upload(merge_memory, inputs=[mem_model, merge_btn], outputs=memory_output)
+
+        clear_btn = gr.Button("Clear Memory")
+        clear_btn.click(clear_memory, inputs=[mem_model], outputs=memory_output)
 
-demo.launch(server_name="0.0.0.0", server_port=7860)
+demo.launch()