Update app.py

app.py

@@ -1,88 +1,31 @@
 #!/usr/bin/env python3
 import os
-import re
+import shutil
 import streamlit as st
-import streamlit.components.v1 as components
-from urllib.parse import quote
 import pandas as pd
 import torch
-import torch.nn as nn
-import torch.optim as optim
-from torch.utils.data import DataLoader, TensorDataset
-import base64
-import glob
-import time
 from transformers import AutoModelForCausalLM, AutoTokenizer
 from torch.utils.data import Dataset, DataLoader
 import csv
+import time
 from dataclasses import dataclass
 from typing import Optional
 
 # Page Configuration
 st.set_page_config(
-    page_title="AI Knowledge Tree Builder 📈🌿",
-    page_icon="🌳✨",
+    page_title="SFT Model Builder 🚀",
+    page_icon="🤖",
     layout="wide",
-    initial_sidebar_state="auto",
+    initial_sidebar_state="expanded",
 )
 
-# Predefined Knowledge Trees
-trees = {
-    "ML Engineering": """
-0. ML Engineering 🌐
-1. Data Preparation
-- Load Data 📊
-- Preprocess Data 🛠️
-2. Model Building
-- Train Model 🤖
-- Evaluate Model 📈
-3. Deployment
-- Deploy Model 🚀
-    """,
-    "Health": """
-0. Health and Wellness 🌿
-1. Physical Health
-- Exercise 🏋️
-- Nutrition 🍎
-2. Mental Health
-- Meditation 🧘
-- Therapy 🛋️
-    """,
-}
-
-# Project Seeds
-project_seeds = {
-    "Code Project": """
-0. Code Project 📂
-1. app.py 🐍
-2. requirements.txt 📦
-3. README.md 📄
-    """,
-    "Papers Project": """
-0. Papers Project 📚
-1. markdown 📝
-2. mermaid 🖼️
-3. huggingface.co 🤗
-    """,
-    "AI Project": """
-0. AI Project 🤖
-1. Streamlit Torch Transformers
-- Streamlit 🌐
-- Torch 🔥
-- Transformers 🤖
-2. SFT Fine-Tuning
-- SFT 🤓
-- Small Models 📉
-    """,
-}
-
 # Meta class for model configuration
 class ModelMeta(type):
     def __new__(cls, name, bases, attrs):
         attrs['registry'] = {}
         return super().__new__(cls, name, bases, attrs)
 
-# Base Model Configuration Class
+# Model Configuration Class
 @dataclass
 class ModelConfig(metaclass=ModelMeta):
     name: str
@@ -121,10 +64,10 @@ class SFTDataset(Dataset):
         return {
             "input_ids": encoding["input_ids"].squeeze(),
             "attention_mask": encoding["attention_mask"].squeeze(),
-            "labels": encoding["input_ids"].squeeze()  # For causal LM, labels are the same as input_ids
+            "labels": encoding["input_ids"].squeeze()
         }
 
-# Model Builder Class with SFT and Evaluation
+# Model Builder Class
 class ModelBuilder:
     def __init__(self):
         self.config = None
@@ -132,62 +75,53 @@ class ModelBuilder:
         self.tokenizer = None
         self.sft_data = None
 
-    def load_base_model(self, model_name: str):
-        """Load base model from Hugging Face"""
-        with st.spinner("Loading base model..."):
-            self.model = AutoModelForCausalLM.from_pretrained(model_name)
-            self.tokenizer = AutoTokenizer.from_pretrained(model_name)
+    def load_model(self, model_path: str):
+        """Load a model from a path"""
+        with st.spinner("Loading model... ⏳"):
+            self.model = AutoModelForCausalLM.from_pretrained(model_path)
+            self.tokenizer = AutoTokenizer.from_pretrained(model_path)
             if self.tokenizer.pad_token is None:
                 self.tokenizer.pad_token = self.tokenizer.eos_token
-        st.success("Base model loaded!")
+        st.success("Model loaded! ✅")
         return self
 
     def fine_tune_sft(self, csv_path: str, epochs: int = 3, batch_size: int = 4):
         """Perform Supervised Fine-Tuning with CSV data"""
-        # Load CSV data
         self.sft_data = []
         with open(csv_path, "r") as f:
             reader = csv.DictReader(f)
             for row in reader:
                 self.sft_data.append({"prompt": row["prompt"], "response": row["response"]})
 
-        # Prepare dataset and dataloader
         dataset = SFTDataset(self.sft_data, self.tokenizer)
         dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)
+        optimizer = torch.optim.AdamW(self.model.parameters(), lr=2e-5)
 
-        # Set up optimizer
-        optimizer = optim.AdamW(self.model.parameters(), lr=2e-5)
-
-        # Training loop
         self.model.train()
         for epoch in range(epochs):
-            with st.spinner(f"Training epoch {epoch + 1}/{epochs}..."):
+            with st.spinner(f"Training epoch {epoch + 1}/{epochs}... ⚙️"):
                 total_loss = 0
                 for batch in dataloader:
                     optimizer.zero_grad()
                     input_ids = batch["input_ids"].to(self.model.device)
                     attention_mask = batch["attention_mask"].to(self.model.device)
                     labels = batch["labels"].to(self.model.device)
-
-                    outputs = self.model(
-                        input_ids=input_ids,
-                        attention_mask=attention_mask,
-                        labels=labels
-                    )
+                    outputs = self.model(input_ids=input_ids, attention_mask=attention_mask, labels=labels)
                     loss = outputs.loss
                     loss.backward()
                     optimizer.step()
                     total_loss += loss.item()
                 st.write(f"Epoch {epoch + 1} completed. Average loss: {total_loss / len(dataloader):.4f}")
-        st.success("SFT Fine-tuning completed!")
+        st.success("SFT Fine-tuning completed! 🎉")
         return self
 
     def save_model(self, path: str):
         """Save the fine-tuned model"""
-        with st.spinner("Saving model..."):
+        with st.spinner("Saving model... 💾"):
+            os.makedirs(os.path.dirname(path), exist_ok=True)
             self.model.save_pretrained(path)
             self.tokenizer.save_pretrained(path)
-        st.success("Model saved!")
+        st.success(f"Model saved at {path}! ✅")
 
     def evaluate(self, prompt: str):
         """Evaluate the model with a prompt"""
@@ -198,295 +132,115 @@ class ModelBuilder:
             return self.tokenizer.decode(outputs[0], skip_special_tokens=True)
 
 # Utility Functions
-def sanitize_label(label):
-    """Remove invalid characters for Mermaid labels."""
-    return re.sub(r'[^\w\s-]', '', label).replace(' ', '_')
-
-def sanitize_filename(label):
-    """Make a valid filename from a label."""
-    return re.sub(r'[^\w\s-]', '', label).replace(' ', '_')
-
-def parse_outline_to_mermaid(outline_text, search_agent):
-    """Convert tree outline to Mermaid syntax with clickable nodes."""
-    lines = outline_text.strip().split('\n')
-    nodes, edges, clicks, stack = [], [], [], []
-    for line in lines:
-        indent = len(line) - len(line.lstrip())
-        level = indent // 4
-        label = re.sub(r'^[#*\->\d\.\s]+', '', line.strip())
-        if label:
-            node_id = f"N{len(nodes)}"
-            sanitized_label = sanitize_label(label)
-            nodes.append(f'{node_id}["{label}"]')
-            search_url = search_urls[search_agent](label)
-            clicks.append(f'click {node_id} "{search_url}" _blank')
-            if stack:
-                parent_level = stack[-1][0]
-                if level > parent_level:
-                    edges.append(f"{stack[-1][1]} --> {node_id}")
-                    stack.append((level, node_id))
-                else:
-                    while stack and stack[-1][0] >= level:
-                        stack.pop()
-                    if stack:
-                        edges.append(f"{stack[-1][1]} --> {node_id}")
-                    stack.append((level, node_id))
-            else:
-                stack.append((level, node_id))
-    return "%%{init: {'themeVariables': {'fontSize': '18px'}}}%%\nflowchart LR\n" + "\n".join(nodes + edges + clicks)
-
-def generate_mermaid_html(mermaid_code):
-    """Generate HTML to display Mermaid diagram."""
-    return f"""
-    <html><head><script src="https://cdn.jsdelivr.net/npm/mermaid/dist/mermaid.min.js"></script>
-    <style>.centered-mermaid{{display:flex;justify-content:center;margin:20px auto;}}</style></head>
-    <body><div class="mermaid centered-mermaid">{mermaid_code}</div>
-    <script>mermaid.initialize({{startOnLoad:true}});</script></body></html>
-    """
-
-def grow_tree(base_tree, new_node_name, parent_node):
-    """Add a new node to the tree under a specified parent."""
-    lines = base_tree.strip().split('\n')
-    new_lines = []
-    added = False
-    for line in lines:
-        new_lines.append(line)
-        if parent_node in line and not added:
-            indent = len(line) - len(line.lstrip())
-            new_lines.append(f"{' ' * (indent + 4)}- {new_node_name} 🌱")
-            added = True
-    return "\n".join(new_lines)
-
-def get_download_link(file_path, mime_type="text/plain"):
+def get_download_link(file_path, mime_type="text/plain", label="Download"):
     """Generate a download link for a file."""
     with open(file_path, 'rb') as f:
         data = f.read()
     b64 = base64.b64encode(data).decode()
-    return f'<a href="data:{mime_type};base64,{b64}" download="{file_path}">Download {file_path}</a>'
-
-def save_tree_to_file(tree_text, parent_node, new_node):
-    """Save tree to a markdown file with name based on nodes."""
-    root_node = tree_text.strip().split('\n')[0].split('.')[1].strip() if tree_text.strip() else "Knowledge_Tree"
-    filename = f"{sanitize_filename(root_node)}_{sanitize_filename(parent_node)}_{sanitize_filename(new_node)}_{int(time.time())}.md"
-    
-    mermaid_code = parse_outline_to_mermaid(tree_text, "🔮Google")  # Default search engine for saved trees
-    export_md = f"# Knowledge Tree: {root_node}\n\n## Outline\n{tree_text}\n\n## Mermaid Diagram\n```mermaid\n{mermaid_code}\n```"
-    
-    with open(filename, "w") as f:
-        f.write(export_md)
-    return filename
-
-def load_trees_from_files():
-    """Load all saved tree markdown files."""
-    tree_files = glob.glob("*.md")
-    trees_dict = {}
-    
-    for file in tree_files:
-        if file != "README.md" and file != "knowledge_tree.md":  # Skip project README and temp export
-            try:
-                with open(file, 'r') as f:
-                    content = f.read()
-                    # Extract the tree name from the first line
-                    match = re.search(r'# Knowledge Tree: (.*)', content)
-                    if match:
-                        tree_name = match.group(1)
-                    else:
-                        tree_name = os.path.splitext(file)[0]
-                    
-                    # Extract the outline section
-                    outline_match = re.search(r'## Outline\n(.*?)(?=\n## |$)', content, re.DOTALL)
-                    if outline_match:
-                        tree_outline = outline_match.group(1).strip()
-                        trees_dict[f"{tree_name} ({file})"] = tree_outline
-            except Exception as e:
-                print(f"Error loading {file}: {e}")
-    
-    return trees_dict
+    return f'<a href="data:{mime_type};base64,{b64}" download="{os.path.basename(file_path)}">{label} 📥</a>'
 
-# Search Agents (Highest resolution social network default: X)
-search_urls = {
-    "📚📖ArXiv": lambda k: f"/?q={quote(k)}",
-    "🔮Google": lambda k: f"https://www.google.com/search?q={quote(k)}",
-    "📺Youtube": lambda k: f"https://www.youtube.com/results?search_query={quote(k)}",
-    "🔭Bing": lambda k: f"https://www.bing.com/search?q={quote(k)}",
-    "💡Truth": lambda k: f"https://truthsocial.com/search?q={quote(k)}",
-    "📱X": lambda k: f"https://twitter.com/search?q={quote(k)}",
-}
+def get_model_files():
+    """List all saved model directories."""
+    return [d for d in glob.glob("models/*") if os.path.isdir(d)]
 
 # Main App
-st.title("🌳 AI Knowledge Tree Builder 🌱")
-
-# Sidebar with saved trees
-st.sidebar.title("Saved Trees")
-saved_trees = load_trees_from_files()
-selected_saved_tree = st.sidebar.selectbox("Select a saved tree", ["None"] + list(saved_trees.keys()))
-
-# Select Project Type
-project_type = st.selectbox("Select Project Type", ["Code Project", "Papers Project", "AI Project"])
-
-# Initialize or load tree
-if 'current_tree' not in st.session_state:
-    if selected_saved_tree != "None" and selected_saved_tree in saved_trees:
-        st.session_state['current_tree'] = saved_trees[selected_saved_tree]
-    else:
-        st.session_state['current_tree'] = trees.get("ML Engineering", project_seeds[project_type])
-elif selected_saved_tree != "None" and selected_saved_tree in saved_trees:
-    st.session_state['current_tree'] = saved_trees[selected_saved_tree]
-
-# Select Search Agent for Node Links
-search_agent = st.selectbox("Select Search Agent for Node Links", list(search_urls.keys()), index=5)  # Default to X
-
-# Tree Growth
-new_node = st.text_input("Add New Node")
-parent_node = st.text_input("Parent Node")
-if st.button("Grow Tree 🌱") and new_node and parent_node:
-    st.session_state['current_tree'] = grow_tree(st.session_state['current_tree'], new_node, parent_node)
-    
-    # Save to a new file with the node names
-    saved_file = save_tree_to_file(st.session_state['current_tree'], parent_node, new_node)
-    st.success(f"Added '{new_node}' under '{parent_node}' and saved to {saved_file}!")
-    
-    # Also update the temporary current_tree.md for compatibility
-    with open("current_tree.md", "w") as f:
-        f.write(st.session_state['current_tree'])
+st.title("SFT Model Builder 🤖🚀")
+
+# Sidebar for Model Management
+st.sidebar.header("Model Management 🗂️")
+model_dirs = get_model_files()
+selected_model = st.sidebar.selectbox("Select Saved Model", ["None"] + model_dirs)
+
+if selected_model != "None" and st.sidebar.button("Load Model 📂"):
+    if 'builder' not in st.session_state:
+        st.session_state['builder'] = ModelBuilder()
+    st.session_state['builder'].load_model(selected_model)
+    st.session_state['model_loaded'] = True
     st.rerun()
 
-# Display Mermaid Diagram
-st.markdown("### Knowledge Tree Visualization")
-mermaid_code = parse_outline_to_mermaid(st.session_state['current_tree'], search_agent)
-components.html(generate_mermaid_html(mermaid_code), height=600)
-
-# Export Tree
-if st.button("Export Tree as Markdown"):
-    export_md = f"# Knowledge Tree\n\n## Outline\n{st.session_state['current_tree']}\n\n## Mermaid Diagram\n```mermaid\n{mermaid_code}\n```"
-    with open("knowledge_tree.md", "w") as f:
-        f.write(export_md)
-    st.markdown(get_download_link("knowledge_tree.md", "text/markdown"), unsafe_allow_html=True)
-
-# AI Project: Model Building Options
-if project_type == "AI Project":
-    st.subheader("AI Model Building Options")
-    model_option = st.radio("Choose Model Building Method", ["Minimal ML Model from CSV", "SFT Fine-Tuning"])
-
-    if model_option == "Minimal ML Model from CSV":
-        st.write("### Build Minimal ML Model from CSV")
-        uploaded_file = st.file_uploader("Upload CSV", type="csv")
-        if uploaded_file:
-            df = pd.read_csv(uploaded_file)
-            st.write("Columns:", df.columns.tolist())
-            feature_cols = st.multiselect("Select feature columns", df.columns)
-            target_col = st.selectbox("Select target column", df.columns)
-            if st.button("Train Model"):
-                X = df[feature_cols].values
-                y = df[target_col].values
-                X_tensor = torch.tensor(X, dtype=torch.float32)
-                y_tensor = torch.tensor(y, dtype=torch.float32).view(-1, 1)
-                dataset = TensorDataset(X_tensor, y_tensor)
-                loader = DataLoader(dataset, batch_size=32, shuffle=True)
-                model = nn.Linear(X.shape[1], 1)
-                criterion = nn.MSELoss()
-                optimizer = optim.Adam(model.parameters(), lr=0.01)
-                for epoch in range(10):
-                    for batch_X, batch_y in loader:
-                        optimizer.zero_grad()
-                        outputs = model(batch_X)
-                        loss = criterion(outputs, batch_y)
-                        loss.backward()
-                        optimizer.step()
-                torch.save(model.state_dict(), "model.pth")
-                app_code = f"""
-import streamlit as st
-import torch
-import torch.nn as nn
-
-model = nn.Linear({len(feature_cols)}, 1)
-model.load_state_dict(torch.load("model.pth"))
-model.eval()
-
-st.title("ML Model Demo")
-inputs = []
-for col in {feature_cols}:
-    inputs.append(st.number_input(col))
-if st.button("Predict"):
-    input_tensor = torch.tensor([inputs], dtype=torch.float32)
-    prediction = model(input_tensor).item()
-    st.write(f"Predicted {target_col}: {{prediction}}")
-"""
-                with open("app.py", "w") as f:
-                    f.write(app_code)
-                reqs = "streamlit\ntorch\npandas\n"
-                with open("requirements.txt", "w") as f:
-                    f.write(reqs)
-                readme = """
-# ML Model Demo
-
-## How to run
-1. Install requirements: `pip install -r requirements.txt`
-2. Run the app: `streamlit run app.py`
-3. Input feature values and click "Predict".
-"""
-                with open("README.md", "w") as f:
-                    f.write(readme)
-                st.markdown(get_download_link("model.pth", "application/octet-stream"), unsafe_allow_html=True)
-                st.markdown(get_download_link("app.py", "text/plain"), unsafe_allow_html=True)
-                st.markdown(get_download_link("requirements.txt", "text/plain"), unsafe_allow_html=True)
-                st.markdown(get_download_link("README.md", "text/markdown"), unsafe_allow_html=True)
-
-    elif model_option == "SFT Fine-Tuning":
-        st.write("### SFT Fine-Tuning with Small Models")
-        
-        # Model Configuration
-        with st.expander("Model Configuration", expanded=True):
-            base_model = st.selectbox(
-                "Select Base Model",
-                ["distilgpt2", "gpt2", "EleutherAI/pythia-70m"],
-                help="Choose a small model for fine-tuning"
-            )
-            model_name = st.text_input("Model Name", f"sft-model-{int(time.time())}")
-            domain = st.text_input("Target Domain", "general")
-
-        # Initialize ModelBuilder
-        if 'builder' not in st.session_state:
-            st.session_state['builder'] = ModelBuilder()
-
-        # Load Sample Model
-        if st.button("Load Sample Model"):
-            st.session_state['builder'].load_base_model(base_model)
-            st.session_state['model_loaded'] = True
-            st.rerun()
-
-        # Generate and Export Sample CSV
-        if st.button("Generate Sample CSV"):
+# Main UI with Tabs
+tab1, tab2, tab3 = st.tabs(["Build New Model 🌱", "Fine-Tune Model 🔧", "Test Model 🧪"])
+
+with tab1:
+    st.header("Build New Model 🌱")
+    base_model = st.selectbox(
+        "Select Base Model",
+        ["distilgpt2", "gpt2", "EleutherAI/pythia-70m"],
+        help="Choose a small model to start with"
+    )
+    model_name = st.text_input("Model Name", f"new-model-{int(time.time())}")
+    domain = st.text_input("Target Domain", "general")
+
+    if st.button("Download Model ⬇️"):
+        config = ModelConfig(name=model_name, base_model=base_model, size="small", domain=domain)
+        builder = ModelBuilder()
+        builder.load_model(base_model)
+        builder.save_model(config.model_path)
+        st.session_state['builder'] = builder
+        st.session_state['model_loaded'] = True
+        st.success(f"Model downloaded and saved to {config.model_path}! 🎉")
+        st.rerun()
+
+with tab2:
+    st.header("Fine-Tune Model 🔧")
+    if 'builder' not in st.session_state or not st.session_state.get('model_loaded', False):
+        st.warning("Please download or load a model first! ⚠️")
+    else:
+        # Generate Sample CSV
+        if st.button("Generate Sample CSV 📝"):
             sample_data = [
                 {"prompt": "What is AI?", "response": "AI is artificial intelligence, simulating human intelligence in machines."},
                 {"prompt": "Explain machine learning", "response": "Machine learning is a subset of AI where models learn from data."},
                 {"prompt": "What is a neural network?", "response": "A neural network is a model inspired by the human brain."},
             ]
-            with open("sft_data.csv", "w", newline="") as f:
+            csv_path = f"sft_data_{int(time.time())}.csv"
+            with open(csv_path, "w", newline="") as f:
                 writer = csv.DictWriter(f, fieldnames=["prompt", "response"])
                 writer.writeheader()
                 writer.writerows(sample_data)
-            st.markdown(get_download_link("sft_data.csv", "text/csv"), unsafe_allow_html=True)
-            st.success("Sample CSV generated as 'sft_data.csv'!")
+            st.markdown(get_download_link(csv_path, "text/csv", "Download Sample CSV"), unsafe_allow_html=True)
+            st.success(f"Sample CSV generated as {csv_path}! ✅")
 
         # Upload CSV and Fine-Tune
-        uploaded_csv = st.file_uploader("Upload CSV for SFT (or use generated sample)", type="csv")
-        if st.button("Fine-Tune Model") and (uploaded_csv or os.path.exists("sft_data.csv")):
-            if not hasattr(st.session_state['builder'], 'model') or st.session_state['builder'].model is None:
-                st.session_state['builder'].load_base_model(base_model)
-            
-            csv_path = "sft_data.csv"
-            if uploaded_csv:
-                with open(csv_path, "wb") as f:
-                    f.write(uploaded_csv.read())
-            
-            with st.status("Fine-tuning model...", expanded=True) as status:
+        uploaded_csv = st.file_uploader("Upload CSV for SFT", type="csv")
+        if uploaded_csv and st.button("Fine-Tune with Uploaded CSV 🔄"):
+            csv_path = f"uploaded_sft_data_{int(time.time())}.csv"
+            with open(csv_path, "wb") as f:
+                f.write(uploaded_csv.read())
+            new_model_name = f"{st.session_state['builder'].config.name}-sft-{int(time.time())}"
+            new_config = ModelConfig(name=new_model_name, base_model=st.session_state['builder'].config.base_model, size="small", domain=st.session_state['builder'].config.domain)
+            st.session_state['builder'].config = new_config
+            with st.status("Fine-tuning model... ⏳", expanded=True) as status:
                 st.session_state['builder'].fine_tune_sft(csv_path)
-                st.session_state['builder'].save_model(st.session_state['builder'].config.model_path)
-                status.update(label="Model fine-tuning completed!", state="complete")
-            
-            # Generate deployment files
-            config = ModelConfig(name=model_name, base_model=base_model, size="small", domain=domain)
+                st.session_state['builder'].save_model(new_config.model_path)
+                status.update(label="Fine-tuning completed! 🎉", state="complete")
+            st.markdown(get_download_link(f"{new_config.model_path}/pytorch_model.bin", "application/octet-stream", "Download Fine-Tuned Model"), unsafe_allow_html=True)
+            st.rerun()
+
+with tab3:
+    st.header("Test Model 🧪")
+    if 'builder' not in st.session_state or not st.session_state.get('model_loaded', False):
+        st.warning("Please download or load a model first! ⚠️")
+    else:
+        if st.session_state['builder'].sft_data:
+            st.write("Testing with SFT Data:")
+            for item in st.session_state['builder'].sft_data[:3]:
+                prompt = item["prompt"]
+                expected = item["response"]
+                generated = st.session_state['builder'].evaluate(prompt)
+                st.write(f"**Prompt**: {prompt}")
+                st.write(f"**Expected**: {expected}")
+                st.write(f"**Generated**: {generated}")
+                st.write("---")
+
+        test_prompt = st.text_area("Enter Test Prompt", "What is AI?")
+        if st.button("Run Test ▶️"):
+            result = st.session_state['builder'].evaluate(test_prompt)
+            st.write(f"**Generated Response**: {result}")
+
+        # Export Model Files
+        if st.button("Export Model Files 📦"):
+            config = st.session_state['builder'].config
             app_code = f"""
 import streamlit as st
 from transformers import AutoModelForCausalLM, AutoTokenizer
@@ -517,27 +271,7 @@ if st.button("Generate"):
             with open("sft_README.md", "w") as f:
                 f.write(readme)
             
-            st.markdown(get_download_link("sft_app.py", "text/plain"), unsafe_allow_html=True)
-            st.markdown(get_download_link("sft_requirements.txt", "text/plain"), unsafe_allow_html=True)
-            st.markdown(get_download_link("sft_README.md", "text/markdown"), unsafe_allow_html=True)
-            st.write(f"Model saved at: {config.model_path}")
-            st.rerun()
-
-        # Test and Evaluate Model
-        if 'model_loaded' in st.session_state and st.session_state['builder'].model is not None:
-            st.write("### Test and Evaluate Fine-Tuned Model")
-            if st.session_state['builder'].sft_data:
-                st.write("Testing with SFT data:")
-                for item in st.session_state['builder'].sft_data[:3]:  # Show up to 3 examples
-                    prompt = item["prompt"]
-                    expected = item["response"]
-                    generated = st.session_state['builder'].evaluate(prompt)
-                    st.write(f"**Prompt**: {prompt}")
-                    st.write(f"**Expected**: {expected}")
-                    st.write(f"**Generated**: {generated}")
-                    st.write("---")
-            
-            test_prompt = st.text_area("Enter a custom prompt to test", "What is AI?")
-            if st.button("Test Model"):
-                result = st.session_state['builder'].evaluate(test_prompt)
-                st.write(f"**Generated Response**: {result}")
+            st.markdown(get_download_link("sft_app.py", "text/plain", "Download App"), unsafe_allow_html=True)
+            st.markdown(get_download_link("sft_requirements.txt", "text/plain", "Download Requirements"), unsafe_allow_html=True)
+            st.markdown(get_download_link("sft_README.md", "text/markdown", "Download README"), unsafe_allow_html=True)
+            st.success("Model files exported! ✅")