Update app.py

app.py

@@ -1,199 +1,343 @@
+import torch
+from torch.optim import AdamW
+from torch.utils.data import DataLoader
+from transformers import (
+    GPT2Config,
+    AutoModelForCausalLM,
+    DataCollatorForLanguageModeling,
+    PreTrainedTokenizerFast,
+    DataCollatorWithPadding,
+    GenerationConfig,
+)
+from trl import PPOConfig, PPOTrainer, AutoModelForCausalLMWithValueHead
+from tokenizers import Tokenizer
+from tokenizers.models import BPE
+from tokenizers.trainers import BpeTrainer
+from tokenizers.pre_tokenizers import Whitespace
+from datasets import Dataset, load_dataset
+import os
+import torch.nn as nn
 import gradio as gr
 import pandas as pd
-from datasets import load_dataset, get_dataset_split_names
-from huggingface_hub import HfApi
-import os
-import pathlib
-import uuid
-import logging
-
-# --- Setup Logging ---
-logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
-
-# --- Embedding Atlas Imports ---
-from embedding_atlas.data_source import DataSource
-from embedding_atlas.server import make_server
-from embedding_atlas.projection import compute_text_projection
-from embedding_atlas.utils import Hasher
-
-# --- Helper function from embedding_atlas/cli.py ---
-def find_column_name(existing_names, candidate):
-    """Finds a unique column name, appending '_1', '_2', etc. if the candidate name already exists."""
-    if candidate not in existing_names:
-        return candidate
+import time
+import re
+from sentence_transformers import SentenceTransformer
+import faiss
+
+# --- Configuration & Global State ---
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+SFT_MODEL_PATH = "./sft_model"
+RAG_INDEX_PATH = "faiss_index.bin"
+TOKENIZER_FILE_PATH = "savant_tokenizer.json"
+CACHE_DIR = "./hf_cache"
+
+SAVANT_MODEL = None
+SAVANT_TOKENIZER = None
+RAG_DATABASE = None
+full_dataset_for_rag = None
+
+# --- RAG Database Class ---
+class VectorDatabase:
+    def __init__(self, embedder_model_name='all-MiniLM-L6-v2'):
+        self.embedder = SentenceTransformer(embedder_model_name, device=str(DEVICE), cache_folder=CACHE_DIR)
+        self.index = None
+        self.documents = []
+    
+    def build_index(self, texts):
+        print("Building RAG vector index...")
+        self.documents = texts
+        embeddings = self.embedder.encode(texts, convert_to_tensor=True, show_progress_bar=True)
+        self.index = faiss.IndexFlatL2(embeddings.shape[1])
+        self.index.add(embeddings.cpu().numpy())
+        print(f"RAG Index built with {len(self.documents)} documents.")
+    
+    def save_index(self, path):
+        if self.index:
+            faiss.write_index(self.index, path)
+            print(f"RAG Index saved to {path}")
+    
+    def search(self, query, k=3):
+        if self.index is None:
+            return []
+        query_embedding = self.embedder.encode([query], convert_to_tensor=True)
+        distances, indices = self.index.search(query_embedding.cpu().numpy(), k)
+        return [self.documents[i] for i in indices[0]]
+
+# --- Core Logic ---
+def create_tokenizer_file_from_dataset(dataset, save_path=TOKENIZER_FILE_PATH):
+    corpus_path = "temp_corpus.txt"
+    with open(corpus_path, "w", encoding="utf-8") as f:
+        for item in dataset:
+            if item and item.get('question') and item.get('answer'):
+                f.write(str(item['question']) + " " + str(item['answer']) + "\n")
+    
+    raw_tokenizer = Tokenizer(BPE(unk_token="[UNK]"))
+    raw_tokenizer.pre_tokenizer = Whitespace()
+    special_tokens = ["[UNK]", "[PAD]", "<|startoftext|>", "<|endoftext|>"]
+    trainer = BpeTrainer(vocab_size=8192, special_tokens=special_tokens)
+    raw_tokenizer.train(files=[corpus_path], trainer=trainer)
+    os.remove(corpus_path)
+    raw_tokenizer.save(save_path)
+    return save_path
+
+def create_seed_model(config):
+    return AutoModelForCausalLM.from_config(config)
+
+def get_sft_curriculum():
+    return [{"name": "Phase 1: Foundational Math (SFT)", "dataset_name": "openai/gsm8k", "dataset_config": "main", "data_slice": slice(0, 100), "epochs": 1, "learning_rate": 5e-5, "mastery_threshold": 16.0}]
+
+def get_rl_curriculum():
+    return {"name": "Phase 2: Problem Solving (RL)", "dataset_name": "openai/gsm8k", "dataset_config": "main", "data_slice": slice(1000, 1100), "num_steps": 100}
+
+def get_folder_files(folder_path):
+    if not os.path.isdir(folder_path): return []
+    return [os.path.join(folder_path, f) for f in os.listdir(folder_path)]
+
+def extract_answer(text):
+    text = str(text)
+    match = re.search(r'\\boxed\{([^}]*)\}', text)
+    if match:
+        ans = match.group(1).strip().replace(",", "")
+        try: return float(ans)
+        except ValueError: return None
+    matches = re.findall(r'(\d+\.?\d*|\.\d+)', text)
+    if matches:
+        try: return float(matches[-1])
+        except ValueError: return None
+    return None
+
+# --- Master Training Process ---
+def run_sft_phase(artifact_files):
+    global full_dataset_for_rag
+    log_text = "--- Starting Phase 1: Supervised Fine-Tuning (SFT) ---\n"
+    yield log_text, None, gr.Group(visible=False), artifact_files
+    
+    stage = get_sft_curriculum()[0]
+    full_dataset_for_rag = load_dataset(stage['dataset_name'], name=stage['dataset_config'], split='train', cache_dir=CACHE_DIR)
+    tokenizer_file = create_tokenizer_file_from_dataset(full_dataset_for_rag)
+    artifact_files.append(tokenizer_file)
+    
+    wrapped_tokenizer = PreTrainedTokenizerFast(tokenizer_file=tokenizer_file)
+    wrapped_tokenizer.pad_token = "[PAD]"
+    
+    log_text += f"✓ Tokenizer created.\n"
+    yield log_text, None, gr.Group(visible=False), artifact_files
+    
+    config = GPT2Config(vocab_size=wrapped_tokenizer.vocab_size, n_positions=512, n_layer=4, n_head=4, n_embd=256, pad_token_id=wrapped_tokenizer.pad_token_id, eos_token_id=wrapped_tokenizer.eos_token_id)
+    
+    model = create_seed_model(config).to(DEVICE)
+    log_text += f"✓ Seed Model (GPT-2 Style) Initialized.\n\n"
+    yield log_text, None, gr.Group(visible=False), artifact_files
+    
+    data_collator = DataCollatorForLanguageModeling(tokenizer=wrapped_tokenizer, mlm=False)
+    stage_ds = full_dataset_for_rag.select(range(stage['data_slice'].start, stage['data_slice'].stop))
+    
+    tokenized_dataset = stage_ds.map(lambda ex: wrapped_tokenizer([q + " " + a for q, a in zip(ex['question'], ex['answer'])], truncation=True, padding="max_length", max_length=512), batched=True, remove_columns=stage_ds.column_names)
+    
+    dataloader = DataLoader(tokenized_dataset, batch_size=4, shuffle=True, collate_fn=data_collator)
+    optimizer = AdamW(model.parameters(), lr=stage['learning_rate'])
+    
+    loss_history = []
+    avg_epoch_loss = float('inf')
+    
+    for epoch in range(stage['epochs']):
+        base_log_for_epoch = log_text + f"  Starting SFT Epoch {epoch+1}/{stage['epochs']}...\n"
+        yield base_log_for_epoch, None, gr.Group(visible=False), artifact_files
+        for batch_idx, batch in enumerate(dataloader):
+            batch = {k: v.to(DEVICE) for k, v in batch.items()}
+            optimizer.zero_grad()
+            outputs = model(**batch)
+            loss = outputs.loss
+            loss.backward()
+            optimizer.step()
+            if (batch_idx + 1) % 20 == 0:
+                yield base_log_for_epoch + f"    - Batch {batch_idx+1}/{len(dataloader)}\n", None, gr.Group(visible=False), artifact_files
+        avg_epoch_loss = loss.item()
+        loss_history.append({"Phase": "SFT", "Epoch": epoch, "Loss": avg_epoch_loss})
+        loss_df = pd.DataFrame(loss_history)
+        log_text += f"  Epoch {epoch+1}/{stage['epochs']} complete. Loss: {avg_epoch_loss:.4f}\n"
+        yield log_text, gr.LinePlot(loss_df, x="Epoch", y="Loss", color="Phase"), gr.Group(visible=False), artifact_files
+    
+    log_text += f"✓ SFT Phase Complete. Final Loss: {avg_epoch_loss:.4f}\n"
+    if avg_epoch_loss < stage['mastery_threshold']:
+        log_text += f"✓ SFT Mastery Gate PASSED! Saving model...\n\n"
+        model.save_pretrained(SFT_MODEL_PATH)
+        wrapped_tokenizer.save_pretrained(SFT_MODEL_PATH)
+        artifact_files = [f for f in artifact_files if SFT_MODEL_PATH not in os.path.dirname(f)]
+        artifact_files.extend(get_folder_files(SFT_MODEL_PATH))
+        yield log_text, None, gr.Group(visible=False), artifact_files
+        return model, wrapped_tokenizer, log_text, artifact_files
     else:
-        index = 1
-        while True:
-            s = f"{candidate}_{index}"
-            if s not in existing_names:
-                return s
-            index += 1
-
-# --- Hugging Face API Helpers for Dynamic UI ---
-hf_api = HfApi()
-
-def get_user_datasets(username: str):
-    """Fetches all public datasets for a given username or organization."""
-    logging.info(f"Fetching datasets for user: {username}")
-    if not username:
-        return gr.update(choices=[], value=None, interactive=False)
-    try:
-        datasets = hf_api.list_datasets(author=username, full=True)
-        dataset_ids = [d.id for d in datasets if not d.private]
-        logging.info(f"Found {len(dataset_ids)} datasets for {username}.")
-        return gr.update(choices=sorted(dataset_ids), value=None, interactive=True)
-    except Exception as e:
-        logging.error(f"Failed to fetch datasets for {username}: {e}")
-        gr.Warning(f"Could not fetch datasets for user '{username}'.")
-        return gr.update(choices=[], value=None, interactive=False)
-
-def get_dataset_splits(dataset_id: str):
-    """Gets all available splits for a selected dataset."""
-    logging.info(f"Fetching splits for dataset: {dataset_id}")
-    if not dataset_id:
-        return gr.update(choices=[], value=None, interactive=False)
-    try:
-        splits = get_dataset_split_names(dataset_id)
-        logging.info(f"Found splits for {dataset_id}: {splits}")
-        return gr.update(choices=splits, value=splits[0] if splits else None, interactive=True)
-    except Exception as e:
-        logging.error(f"Failed to fetch splits for {dataset_id}: {e}")
-        gr.Warning(f"Could not fetch splits for dataset '{dataset_id}'.")
-        return gr.update(choices=[], value=None, interactive=False)
-
-def get_split_columns(dataset_id: str, split: str):
-    """Gets all columns for a selected split by loading one row."""
-    logging.info(f"Fetching columns for: {dataset_id}, split: {split}")
-    if not dataset_id or not split:
-        return gr.update(choices=[], value=None, interactive=False)
-    try:
-        dataset_sample = load_dataset(dataset_id, split=split, streaming=True)
-        first_row = next(iter(dataset_sample))
-        columns = list(first_row.keys())
-        logging.info(f"Found columns: {columns}")
-        
-        preferred_cols = ['text', 'content', 'instruction', 'question', 'document', 'prompt']
-        best_col = next((col for col in preferred_cols if col in columns), columns[0] if columns else None)
-        logging.info(f"Best default column chosen: {best_col}")
+        log_text += f"✗ SFT Mastery Gate FAILED. Stopping.\n"
+        yield log_text, None, gr.Group(visible=False), artifact_files
+        return None, None, log_text, artifact_files
+
+def run_rl_phase(sft_model, tokenizer, initial_log_text, rl_dataset_slice, artifact_files):
+    global SAVANT_MODEL, SAVANT_TOKENIZER
+    log_text = initial_log_text
+    rl_curriculum = get_rl_curriculum()
+    
+    log_text += f"--- Starting Phase 2: Reinforcement Learning (RL) ---\n"
+    yield log_text, None, gr.Group(visible=False), artifact_files
+    
+    def tokenize_query(examples):
+        return tokenizer(examples["question"], truncation=True, max_length=512)
+
+    rl_dataset = rl_dataset_slice.map(tokenize_query, batched=True)
+
+    ppo_config = PPOConfig(
+        learning_rate=1.41e-5,
+        batch_size=4,
+        mini_batch_size=4,
+        ppo_epochs=4,
+    )
+
+    model_with_value_head = AutoModelForCausalLMWithValueHead.from_pretrained(SFT_MODEL_PATH).to(DEVICE)
+    
+    data_collator = DataCollatorWithPadding(tokenizer)
+    ppo_trainer = PPOTrainer(
+        config=ppo_config,
+        model=model_with_value_head,
+        ref_model=None,
+        tokenizer=tokenizer,
+        dataset=rl_dataset,
+        data_collator=data_collator,
+    )
+
+    # FIXED: Create our own DataLoader to ensure the 'answer' column is preserved in the batch.
+    # The PPOTrainer's internal dataloader drops all columns except for model inputs.
+    dataloader = DataLoader(
+        rl_dataset,
+        batch_size=ppo_config.batch_size,
+        collate_fn=data_collator
+    )
+
+    generation_kwargs = {"max_new_tokens": 100, "pad_token_id": tokenizer.pad_token_id, "do_sample": True, "temperature": 0.7}
+    log_text += "✓ PPOTrainer instantiated. Starting PPO training loop...\n"
+    yield log_text, None, gr.Group(visible=False), artifact_files
+    
+    # FIXED: Iterate over our custom dataloader, not ppo_trainer.dataloader
+    for step, batch in enumerate(dataloader):
+        if step >= rl_curriculum["num_steps"]:
+            break
+
+        query_tensors = batch["input_ids"].to(DEVICE)
+        # This will now work correctly because our dataloader preserves the 'answer' column.
+        ground_truth_answers = batch["answer"]
+
+        response_tensors = ppo_trainer.generate(query_tensors, return_prompt=False, **generation_kwargs)
+        response_texts = [tokenizer.decode(r, skip_special_tokens=True) for r in response_tensors]
         
-        return gr.update(choices=columns, value=best_col, interactive=True)
-    except Exception as e:
-        logging.error(f"Failed to get columns for {dataset_id}/{split}: {e}", exc_info=True)
-        gr.Warning(f"Could not fetch columns for split '{split}'. Error: {e}")
-        return gr.update(choices=[], value=None, interactive=False)
-
-# --- Main Atlas Generation Logic ---
-def generate_atlas(
-    dataset_name: str,
-    split: str,
-    text_column: str,
-    sample_size: int,
-    model_name: str,
-    umap_neighbors: int,
-    umap_min_dist: float,
-    progress=gr.Progress(track_tqdm=True)
-):
-    """
-    Loads data, computes embeddings, and serves the Embedding Atlas UI.
-    """
-    if not all([dataset_name, split, text_column]):
-        raise gr.Error("Please ensure a Dataset, Split, and Text Column are selected.")
+        rewards = []
+        for resp_text, gt_answer in zip(response_texts, ground_truth_answers):
+            pred = extract_answer(resp_text)
+            true = extract_answer(gt_answer)
+            reward_value = 1.0 if pred is not None and true is not None and abs(pred - true) < 1e-2 else 0.0
+            rewards.append(torch.tensor(reward_value, device=DEVICE))
         
-    progress(0, desc=f"Loading dataset '{dataset_name}' [{split}]...")
-    try:
-        dataset = load_dataset(dataset_name, split=split)
-        df = dataset.to_pandas()
-    except Exception as e:
-        raise gr.Error(f"Failed to load data. Error: {e}")
-
-    if sample_size > 0 and sample_size < len(df):
-        progress(0.1, desc=f"Sampling {sample_size} rows...")
-        df = df.sample(n=sample_size, random_state=42).reset_index(drop=True)
-    
-    if text_column not in df.columns:
-        raise gr.Error(f"Column '{text_column}' not found. Please select a valid column.")
-
-    progress(0.2, desc="Computing embeddings and UMAP...")
-    
-    x_col = find_column_name(df.columns, "projection_x")
-    y_col = find_column_name(df.columns, "projection_y")
-    neighbors_col = find_column_name(df.columns, "__neighbors")
-    
-    try:
-        compute_text_projection(
-            df, text_column, x=x_col, y=y_col, neighbors=neighbors_col, model=model_name,
-            umap_args={"n_neighbors": umap_neighbors, "min_dist": umap_min_dist, "metric": "cosine", "random_state": 42},
-        )
-    except Exception as e:
-        raise gr.Error(f"Failed to compute embeddings. Check model name or sample size. Error: {e}")
-
-    progress(0.8, desc="Preparing Atlas data source...")
-    id_col = find_column_name(df.columns, "_row_index")
-    df[id_col] = range(df.shape[0])
-
-    metadata = {"columns": {"id": id_col, "text": text_column, "embedding": {"x": x_col, "y": y_col}, "neighbors": neighbors_col}}
-    hasher = Hasher()
-    hasher.update(f"{dataset_name}-{split}-{text_column}-{sample_size}-{model_name}")
-    identifier = hasher.hexdigest()
-    atlas_dataset = DataSource(identifier, df, metadata)
-    
-    progress(0.9, desc="Mounting visualization UI...")
-    static_path = str((pathlib.Path(__import__('embedding_atlas').__file__).parent / "static").resolve())
-    mount_path = f"/{uuid.uuid4().hex}"
-    atlas_app = make_server(atlas_dataset, static_path=static_path, duckdb_uri="wasm")
-    
-    # --- THIS IS THE FINAL, CORRECT METHOD ---
-    # The 'app' object from gr.Blocks() is a FastAPI app. We use its standard 'mount' method.
-    app.mount(mount_path, atlas_app)
-
-    progress(1.0, desc="Done!")
-    iframe_html = f"<iframe src='{mount_path}' width='100%' height='800px' frameborder='0'></iframe>"
-    return gr.HTML(iframe_html)
-
-# --- Gradio UI Definition ---
-with gr.Blocks(theme=gr.themes.Soft(), title="Embedding Atlas Explorer") as app:
-    gr.Markdown("# Embedding Atlas Explorer")
-    gr.Markdown("Interactively select and visualize any text-based dataset from the Hugging Face Hub.")
+        stats = ppo_trainer.step(query_tensors, response_tensors, rewards)
+
+        if step % 10 == 0:
+            mean_reward = stats.get("ppo/returns/mean", torch.tensor(0.0)).item()
+            log_text += f"Step {step}/{rl_curriculum['num_steps']}, Mean Reward: {mean_reward:.2f}\n"
+            yield log_text, None, gr.Group(visible=False), artifact_files
 
+    log_text += "\n✓ Reinforcement Learning Phase Complete! Saving final policy model...\n"
+    
+    ppo_trainer.save_model(SFT_MODEL_PATH)
+    
+    SAVANT_MODEL = ppo_trainer.model
+    SAVANT_TOKENIZER = tokenizer
+    
+    artifact_files = [f for f in artifact_files if SFT_MODEL_PATH not in os.path.dirname(f)]
+    artifact_files.extend(get_folder_files(SFT_MODEL_PATH))
+    
+    log_text += f"✓ Final RL-tuned model saved to {SFT_MODEL_PATH}\n"
+    yield log_text, None, gr.Group(visible=False), artifact_files
+    
+    return log_text, artifact_files
+
+def run_full_training():
+    global RAG_DATABASE, SAVANT_MODEL, SAVANT_TOKENIZER, full_dataset_for_rag
+    artifact_files = []
+    
+    sft_model, tokenizer, log_text, artifact_files = yield from run_sft_phase(artifact_files)
+    
+    if sft_model is None:
+        yield log_text, None, gr.Group(visible=False), artifact_files
+        return
+    
+    log_text += "--- Building RAG Knowledge Base ---\n"
+    yield log_text, None, gr.Group(visible=False), artifact_files
+    
+    rag_texts = [f"Question: {q}\nAnswer: {a}" for q, a in zip(full_dataset_for_rag['question'], full_dataset_for_rag['answer'])]
+    RAG_DATABASE = VectorDatabase()
+    RAG_DATABASE.build_index(rag_texts)
+    RAG_DATABASE.save_index(RAG_INDEX_PATH)
+    artifact_files.append(RAG_INDEX_PATH)
+    
+    log_text += "✓ RAG Knowledge Base Built.\n\n"
+    yield log_text, None, gr.Group(visible=False), artifact_files
+    
+    rl_curriculum = get_rl_curriculum()
+    rl_dataset_slice = full_dataset_for_rag.select(range(*rl_curriculum['data_slice'].indices(len(full_dataset_for_rag))))
+    
+    final_log_text, final_artifact_files = yield from run_rl_phase(sft_model, tokenizer, log_text, rl_dataset_slice, artifact_files)
+    
+    final_log_text += "\n--- Curriculum Completed: RAG-Powered Savant Ready! ---\n"
+    
+    SAVANT_MODEL = AutoModelForCausalLM.from_pretrained(SFT_MODEL_PATH).to(DEVICE)
+    SAVANT_TOKENIZER = PreTrainedTokenizerFast.from_pretrained(SFT_MODEL_PATH)
+
+    yield final_log_text, None, gr.Group(visible=True), final_artifact_files
+
+def run_savant_inference(user_question):
+    if SAVANT_MODEL is None or SAVANT_TOKENIZER is None or RAG_DATABASE is None:
+        return "Model or RAG Database not ready. Please run the training first."
+    
+    retrieved_context = RAG_DATABASE.search(user_question, k=3)
+    context_str = "\n".join(retrieved_context)
+    augmented_prompt = f"Context:\n{context_str}\n\nQuestion:\n{user_question}\n\nAnswer:"
+    
+    inputs = SAVANT_TOKENIZER(augmented_prompt, return_tensors="pt").to(DEVICE)
+    
+    with torch.no_grad():
+        response_ids = SAVANT_MODEL.generate(**inputs, max_new_tokens=150, pad_token_id=SAVANT_TOKENIZER.pad_token_id, do_sample=True, temperature=0.7, top_p=0.9)
+    
+    response_text = SAVANT_TOKENIZER.decode(response_ids[0], skip_special_tokens=True)
+    
+    if "Answer:" in response_text:
+        final_answer = response_text.split("Answer:")[-1].strip()
+    else:
+        final_answer = response_text
+    
+    return f"**Savant's Answer:**\n{final_answer}\n\n**Context Used by RAG:**\n- {retrieved_context[0]}\n- {retrieved_context[1]}\n- {retrieved_context[2]}"
+
+# --- GRADIO UI DEFINITION ---
+with gr.Blocks(theme=gr.themes.Soft(), title="Savant-Garde Dashboard") as demo:
+    gr.Markdown("# The Savant-Garde: RAG + SFT + RL Savant Factory")
+    
     with gr.Row():
         with gr.Column(scale=1):
-            gr.Markdown("### 1. Select Data")
-            hf_user_input = gr.Textbox(label="Hugging Face User or Org Name", value="Trendyol", placeholder="e.g., 'gradio' or 'google'")
-            dataset_input = gr.Dropdown(label="Select a Dataset", interactive=False)
-            split_input = gr.Dropdown(label="Select a Split", interactive=False)
-            text_column_input = gr.Dropdown(label="Select a Text Column", interactive=False)
-            
-            gr.Markdown("### 2. Configure Visualization")
-            sample_size_input = gr.Slider(label="Number of Samples", minimum=0, maximum=10000, value=2000, step=100)
-            
-            with gr.Accordion("Advanced Settings", open=False):
-                model_input = gr.Dropdown(label="Embedding Model", choices=["all-MiniLM-L6-v2", "all-mpnet-base-v2", "multi-qa-MiniLM-L6-cos-v1"], value="all-MiniLM-L6-v2")
-                umap_neighbors_input = gr.Slider(label="UMAP Neighbors", minimum=2, maximum=100, value=15, step=1, info="Controls local vs. global structure.")
-                umap_min_dist_input = gr.Slider(label="UMAP Min Distance", minimum=0.0, maximum=0.99, value=0.1, step=0.01, info="Controls how tightly points are packed.")
-
-            generate_button = gr.Button("Generate Atlas", variant="primary")
-
-        with gr.Column(scale=3):
-            gr.Markdown("### 3. Explore Atlas")
-            output_html = gr.HTML("<div style='display:flex; justify-content:center; align-items:center; height:800px; border: 1px solid #ddd; border-radius: 5px;'><p>Atlas will be displayed here after generation.</p></div>")
-
-    # --- Chained Event Listeners for Dynamic UI (Corrected Logic) ---
-    hf_user_input.submit(fn=get_user_datasets, inputs=hf_user_input, outputs=dataset_input)
-    
-    dataset_input.change(fn=get_dataset_splits, inputs=dataset_input, outputs=split_input)
-    
-    # The columns are populated only AFTER a split is chosen.
-    split_input.change(fn=get_split_columns, inputs=[dataset_input, split_input], outputs=text_column_input)
-    
-    # --- Button Click Event ---
-    generate_button.click(
-        fn=generate_atlas,
-        inputs=[dataset_input, split_input, text_column_input, sample_size_input, model_input, umap_neighbors_input, umap_min_dist_input],
-        outputs=[output_html],
-    )
+            start_button = gr.Button("🚀 Begin Full Training", variant="primary")
+            log_output = gr.Textbox(label="Live Training Log", interactive=False, lines=25, max_lines=25)
+        
+        with gr.Column(scale=1):
+            gr.Markdown("### Dashboard Visualizations")
+            loss_plot = gr.LinePlot(show_label=False)
+            with gr.Accordion("Downloadable Artifacts", open=True):
+                file_output = gr.File(label="Generated Files", file_count="multiple", interactive=False)
+    
+    with gr.Group(visible=False) as interaction_group:
+        gr.Markdown("---")
+        gr.Markdown("### 🧠 Interact with the RAG-Powered Savant")
+        with gr.Row():
+            question_input = gr.Textbox(label="Ask a Math Question", placeholder="e.g., Janet has 2 apples...")
+            ask_button = gr.Button("Ask Savant")
+        savant_answer = gr.Markdown()
     
-    # Load initial example data on app load
-    app.load(fn=get_user_datasets, inputs=hf_user_input, outputs=dataset_input)
+    start_button.click(fn=run_full_training, inputs=None, outputs=[log_output, loss_plot, interaction_group, file_output])
+    ask_button.click(fn=run_savant_inference, inputs=question_input, outputs=savant_answer)
 
 if __name__ == "__main__":
-    app.launch(debug=True)
+    demo.launch(debug=True)