Update app.py

app.py

@@ -13,17 +13,10 @@ import base64
 import glob
 import time
 from transformers import AutoModelForCausalLM, AutoTokenizer
-from mergekit.config import MergeConfiguration
-from mergekit.merge import Mergekit
-from spectrum import SpectrumAnalyzer
-import yaml
+from torch.utils.data import Dataset, DataLoader
+import csv
 from dataclasses import dataclass
-from typing import Optional, List
-import logging
-
-# Configure logging
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
+from typing import Optional
 
 # Page Configuration
 st.set_page_config(
@@ -77,13 +70,9 @@ project_seeds = {
 - Streamlit 🌐
 - Torch 🔥
 - Transformers 🤖
-2. MergeKit Spectrum
-- MergeKit 🔄
-- Spectrum 📊
-3. Transformers Diffusers Datasets
-- Transformers 🤖
-- Diffusers 🎨
-- Datasets 📊
+2. SFT Fine-Tuning
+- SFT 🤓
+- Small Models 📉
     """,
 }
 
@@ -108,54 +97,96 @@ class ModelConfig(metaclass=ModelMeta):
     def model_path(self):
         return f"models/{self.name}"
 
-# Decorator for pipeline stages
-def pipeline_stage(func):
-    def wrapper(*args, **kwargs):
-        st.spinner(f"Running {func.__name__}...")
-        result = func(*args, **kwargs)
-        st.success(f"Completed {func.__name__}!")
-        return result
-    return wrapper
+# 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"]
+        input_text = f"{prompt} {response}"
+        encoding = self.tokenizer(
+            input_text,
+            max_length=self.max_length,
+            padding="max_length",
+            truncation=True,
+            return_tensors="pt"
+        )
+        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
+        }
 
-# Model Builder Class
+# Model Builder Class with SFT
 class ModelBuilder:
     def __init__(self):
         self.config = None
         self.model = None
         self.tokenizer = None
 
-    @pipeline_stage
     def load_base_model(self, model_name: str):
         """Load base model from Hugging Face"""
-        self.model = AutoModelForCausalLM.from_pretrained(model_name)
-        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
+        with st.spinner("Loading base model..."):
+            self.model = AutoModelForCausalLM.from_pretrained(model_name)
+            self.tokenizer = AutoTokenizer.from_pretrained(model_name)
+            if self.tokenizer.pad_token is None:
+                self.tokenizer.pad_token = self.tokenizer.eos_token
+        st.success("Base model loaded!")
         return self
 
-    @pipeline_stage
-    def apply_merge(self, models_to_merge: List[str], output_dir: str):
-        """Apply Mergekit for model merging"""
-        merge_config = MergeConfiguration(
-            models=models_to_merge,
-            merge_method="linear",
-            output_dir=output_dir
-        )
-        merger = Mergekit(merge_config)
-        merger.run()
-        self.model = AutoModelForCausalLM.from_pretrained(output_dir)
-        return self
-
-    @pipeline_stage
-    def apply_spectrum(self, domain_data: str):
-        """Apply Spectrum for domain specialization"""
-        analyzer = SpectrumAnalyzer(self.model)
-        analyzer.fit(domain_data)
-        self.model = analyzer.specialized_model
+    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
+        data = []
+        with open(csv_path, "r") as f:
+            reader = csv.DictReader(f)
+            for row in reader:
+                data.append({"prompt": row["prompt"], "response": row["response"]})
+
+        # Prepare dataset and dataloader
+        dataset = SFTDataset(data, self.tokenizer)
+        dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)
+
+        # 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}..."):
+                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
+                    )
+                    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):
-        """Save the final model"""
-        self.model.save_pretrained(path)
-        self.tokenizer.save_pretrained(path)
+        """Save the fine-tuned model"""
+        with st.spinner("Saving model..."):
+            self.model.save_pretrained(path)
+            self.tokenizer.save_pretrained(path)
+        st.success("Model saved!")
 
 # Utility Functions
 def sanitize_label(label):
@@ -325,7 +356,7 @@ if st.button("Export Tree as Markdown"):
 # 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", "Advanced Model Pipeline"])
+    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")
@@ -391,46 +422,55 @@ if st.button("Predict"):
                 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 == "Advanced Model Pipeline":
-        st.write("### Advanced Model Building Pipeline")
+    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",
-                ["mistral-7b", "llama-2-7b", "gpt2-medium"]
+                ["distilgpt2", "gpt2", "EleutherAI/pythia-70m"],  # Small models suitable for SFT
+                help="Choose a small model for fine-tuning"
             )
-            model_name = st.text_input("Model Name", "custom-model")
+            model_name = st.text_input("Model Name", "sft-model")
             domain = st.text_input("Target Domain", "general")
-            use_merging = st.checkbox("Apply Model Merging", False)
-            use_spectrum = st.checkbox("Apply Spectrum Specialization", True)
 
-        # Build Model
-        if st.button("Build Advanced Model"):
+        # 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:
+                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'!")
+
+        # Fine-Tune with SFT
+        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")):
             config = ModelConfig(
                 name=model_name,
                 base_model=base_model,
-                size="7B",
+                size="small",
                 domain=domain
             )
             builder = ModelBuilder()
             
-            with st.status("Building advanced model...", expanded=True) as status:
+            # Load CSV
+            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:
                 builder.load_base_model(config.base_model)
-                
-                if use_merging:
-                    models_to_merge = st.multiselect(
-                        "Select Models to Merge",
-                        ["mistral-7b", "llama-2-7b", "gpt2-medium"]
-                    )
-                    builder.apply_merge(models_to_merge, f"merged_{config.name}")
-                
-                if use_spectrum:
-                    domain_data = st.text_area("Enter domain-specific data", "Sample domain data")
-                    builder.apply_spectrum(domain_data)
-                
+                builder.fine_tune_sft(csv_path)
                 builder.save_model(config.model_path)
-                status.update(label="Advanced model built successfully!", state="complete")
+                status.update(label="Model fine-tuning completed!", state="complete")
             
             # Generate deployment files
             app_code = f"""
@@ -440,32 +480,32 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
 model = AutoModelForCausalLM.from_pretrained("{config.model_path}")
 tokenizer = AutoTokenizer.from_pretrained("{config.model_path}")
 
-st.title("Advanced Model Demo")
-input_text = st.text_area("Enter text")
+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)
+    outputs = model.generate(**inputs, max_new_tokens=50)
     st.write(tokenizer.decode(outputs[0], skip_special_tokens=True))
 """
-            with open("advanced_app.py", "w") as f:
+            with open("sft_app.py", "w") as f:
                 f.write(app_code)
-            reqs = "streamlit\ntorch\ntransformers\nmergekit\nspectrum\n"
-            with open("advanced_requirements.txt", "w") as f:
+            reqs = "streamlit\ntorch\ntransformers\n"
+            with open("sft_requirements.txt", "w") as f:
                 f.write(reqs)
             readme = f"""
-# Advanced Model Demo
+# SFT Model Demo
 
 ## How to run
-1. Install requirements: `pip install -r advanced_requirements.txt`
-2. Run the app: `streamlit run advanced_app.py`
-3. Input text and click "Generate".
+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("advanced_README.md", "w") as f:
+            with open("sft_README.md", "w") as f:
                 f.write(readme)
             
-            st.markdown(get_download_link("advanced_app.py", "text/plain"), unsafe_allow_html=True)
-            st.markdown(get_download_link("advanced_requirements.txt", "text/plain"), unsafe_allow_html=True)
-            st.markdown(get_download_link("advanced_README.md", "text/markdown"), unsafe_allow_html=True)
+            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}")
 
 if __name__ == "__main__":