Create app.py

app.py

@@ -0,0 +1,370 @@
+#!/usr/bin/env python3
+import os
+import shutil
+import glob
+import base64
+import streamlit as st
+import pandas as pd
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from torch.utils.data import Dataset, DataLoader
+import csv
+import time
+from dataclasses import dataclass
+from typing import Optional, Tuple
+import zipfile
+import math
+
+# Page Configuration
+st.set_page_config(
+    page_title="SFT Model Builder 🚀",
+    page_icon="🤖",
+    layout="wide",
+    initial_sidebar_state="expanded",
+)
+
+# Model Configuration Class
+@dataclass
+class ModelConfig:
+    name: str
+    base_model: str
+    size: str
+    domain: Optional[str] = None
+    
+    @property
+    def model_path(self):
+        return f"models/{self.name}"
+
+# Custom Dataset for SFT
+class SFTDataset(Dataset):
+    def __init__(self, data, tokenizer, max_length=128):
+        self.data = data
+        self.tokenizer = tokenizer
+        self.max_length = max_length
+
+    def __len__(self):
+        return len(self.data)
+
+    def __getitem__(self, idx):
+        prompt = self.data[idx]["prompt"]
+        response = self.data[idx]["response"]
+        
+        prompt_encoding = self.tokenizer(
+            prompt,
+            max_length=self.max_length // 2,
+            padding="max_length",
+            truncation=True,
+            return_tensors="pt"
+        )
+        
+        full_text = f"{prompt} {response}"
+        full_encoding = self.tokenizer(
+            full_text,
+            max_length=self.max_length,
+            padding="max_length",
+            truncation=True,
+            return_tensors="pt"
+        )
+        
+        input_ids = prompt_encoding["input_ids"].squeeze()
+        attention_mask = prompt_encoding["attention_mask"].squeeze()
+        labels = full_encoding["input_ids"].squeeze()
+        
+        prompt_len = prompt_encoding["input_ids"].ne(self.tokenizer.pad_token_id).sum().item()
+        labels[:prompt_len] = -100  # Mask prompt in loss
+        
+        return {
+            "input_ids": input_ids,
+            "attention_mask": attention_mask,
+            "labels": labels
+        }
+
+# Model Builder Class
+class ModelBuilder:
+    def __init__(self):
+        self.config = None
+        self.model = None
+        self.tokenizer = None
+        self.sft_data = None
+
+    def load_model(self, model_path: str, config: Optional[ModelConfig] = None):
+        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
+            if config:
+                self.config = config
+        st.success("Model loaded! ✅")
+        return self
+
+    def fine_tune_sft(self, csv_path: str, epochs: int = 3, batch_size: int = 4):
+        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"]})
+
+        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)
+
+        self.model.train()
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.model.to(device)
+        for epoch in range(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(device)
+                    attention_mask = batch["attention_mask"].to(device)
+                    labels = batch["labels"].to(device)
+                    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! 🎉")
+        return self
+
+    def save_model(self, path: str):
+        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(f"Model saved at {path}! ✅")
+
+    def evaluate(self, prompt: str):
+        self.model.eval()
+        with torch.no_grad():
+            inputs = self.tokenizer(prompt, return_tensors="pt", max_length=128, truncation=True).to(self.model.device)
+            outputs = self.model.generate(
+                **inputs,
+                max_new_tokens=50,
+                do_sample=True,
+                top_p=0.95,
+                temperature=0.7
+            )
+            return self.tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+# Utility Functions
+def get_download_link(file_path, mime_type="text/plain", label="Download"):
+    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="{os.path.basename(file_path)}">{label} 📥</a>'
+
+def zip_directory(directory_path, zip_path):
+    with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf:
+        for root, _, files in os.walk(directory_path):
+            for file in files:
+                file_path = os.path.join(root, file)
+                arcname = os.path.relpath(file_path, os.path.dirname(directory_path))
+                zipf.write(file_path, arcname)
+
+def get_model_files():
+    return [d for d in glob.glob("models/*") if os.path.isdir(d)]
+
+# Cargo Travel Time Tool
+def calculate_cargo_travel_time(
+    origin_coords: Tuple[float, float],
+    destination_coords: Tuple[float, float],
+    cruising_speed_kmh: float = 750.0
+) -> float:
+    def to_radians(degrees: float) -> float:
+        return degrees * (math.pi / 180)
+    lat1, lon1 = map(to_radians, origin_coords)
+    lat2, lon2 = map(to_radians, destination_coords)
+    EARTH_RADIUS_KM = 6371.0
+    dlon = lon2 - lon1
+    dlat = lat2 - lat1
+    a = (math.sin(dlat / 2) ** 2 + math.cos(lat1) * math.cos(lat2) * math.sin(dlon / 2) ** 2)
+    c = 2 * math.asin(math.sqrt(a))
+    distance = EARTH_RADIUS_KM * c
+    actual_distance = distance * 1.1
+    flight_time = (actual_distance / cruising_speed_kmh) + 1.0
+    return round(flight_time, 2)
+
+# Main App
+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()
+    config = ModelConfig(name=os.path.basename(selected_model), base_model="unknown", size="small", domain="general")
+    st.session_state['builder'].load_model(selected_model, config)
+    st.session_state['model_loaded'] = True
+    st.rerun()
+
+# Main UI with Tabs
+tab1, tab2, tab3, tab4 = st.tabs(["Build New Model 🌱", "Fine-Tune Model 🔧", "Test Model 🧪", "Agentic RAG Demo 🌐"])
+
+with tab1:
+    st.header("Build New Model 🌱")
+    base_model = st.selectbox(
+        "Select Base Model",
+        [
+            "HuggingFaceTB/SmolLM-135M",  # ~270 MB
+            "HuggingFaceTB/SmolLM-360M",  # ~720 MB
+            "Qwen/Qwen1.5-0.5B-Chat",     # ~1 GB
+            "TinyLlama/TinyLlama-1.1B-Chat-v1.0"  # ~2 GB, slightly over but included
+        ],
+        help="Choose a tiny, open-source model (<1 GB except TinyLlama)"
+    )
+    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, config)
+        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:
+        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."},
+            ]
+            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(csv_path, "text/csv", "Download Sample CSV"), unsafe_allow_html=True)
+            st.success(f"Sample CSV generated as {csv_path}! ✅")
+
+        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(new_config.model_path)
+                status.update(label="Fine-tuning completed! 🎉", state="complete")
+            
+            zip_path = f"{new_config.model_path}.zip"
+            zip_directory(new_config.model_path, zip_path)
+            st.markdown(get_download_link(zip_path, "application/zip", "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}")
+
+        if st.button("Export Model Files 📦"):
+            config = st.session_state['builder'].config
+            app_code = f"""
+import streamlit as st
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+model = AutoModelForCausalLM.from_pretrained("{config.model_path}")
+tokenizer = AutoTokenizer.from_pretrained("{config.model_path}")
+
+st.title("SFT Model Demo")
+input_text = st.text_area("Enter prompt")
+if st.button("Generate"):
+    inputs = tokenizer(input_text, return_tensors="pt")
+    outputs = model.generate(**inputs, max_new_tokens=50, do_sample=True, top_p=0.95, temperature=0.7)
+    st.write(tokenizer.decode(outputs[0], skip_special_tokens=True))
+"""
+            with open("sft_app.py", "w") as f:
+                f.write(app_code)
+            reqs = "streamlit\ntorch\ntransformers\n"
+            with open("sft_requirements.txt", "w") as f:
+                f.write(reqs)
+            readme = f"""
+# SFT Model Demo
+
+## How to run
+1. Install requirements: `pip install -r sft_requirements.txt`
+2. Run the app: `streamlit run sft_app.py`
+3. Input a prompt and click "Generate".
+"""
+            with open("sft_README.md", "w") as f:
+                f.write(readme)
+            
+            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! ✅")
+
+with tab4:
+    st.header("Agentic RAG Demo 🌐")
+    st.write("This demo uses tiny models with Agentic RAG to plan a luxury superhero-themed party, enhancing retrieval with DuckDuckGo.")
+
+    if st.button("Run Agentic RAG Demo 🎉"):
+        try:
+            from smolagents import CodeAgent, DuckDuckGoSearchTool, VisitWebpageTool
+
+            # Load selected tiny model
+            tokenizer = AutoTokenizer.from_pretrained("HuggingFaceTB/SmolLM-135M")
+            model = AutoModelForCausalLM.from_pretrained("HuggingFaceTB/SmolLM-135M")
+
+            # Define Agentic RAG agent
+            agent = CodeAgent(
+                model=model,
+                tokenizer=tokenizer,
+                tools=[DuckDuckGoSearchTool(), VisitWebpageTool(), calculate_cargo_travel_time],
+                additional_authorized_imports=["pandas"],
+                planning_interval=5,
+                verbosity_level=2,
+                max_steps=15,
+            )
+            
+            task = """
+Plan a luxury superhero-themed party at Wayne Manor (42.3601° N, 71.0589° W). Search for the latest superhero party trends using DuckDuckGo,
+refine results to include luxury elements (decorations, entertainment, catering), and calculate cargo travel times from key locations 
+(e.g., New York, LA, London) to Wayne Manor. Synthesize a complete plan and return it as a pandas dataframe with at least 6 entries 
+including locations, travel times, and luxury party ideas.
+"""
+            with st.spinner("Running Agentic RAG system... ⏳"):
+                result = agent.run(task)
+            st.write("Agentic RAG Result:")
+            st.write(result)
+        except ImportError:
+            st.error("Please install required packages: `pip install smolagents pandas`")
+        except Exception as e:
+            st.error(f"Error running demo: {str(e)}")