Update app.py

app.py

@@ -124,12 +124,13 @@ class SFTDataset(Dataset):
             "labels": encoding["input_ids"].squeeze()  # For causal LM, labels are the same as input_ids
         }
 
-# Model Builder Class with SFT
+# Model Builder Class with SFT and Evaluation
 class ModelBuilder:
     def __init__(self):
         self.config = None
         self.model = None
         self.tokenizer = None
+        self.sft_data = None
 
     def load_base_model(self, model_name: str):
         """Load base model from Hugging Face"""
@@ -144,14 +145,14 @@ class ModelBuilder:
     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 = []
+        self.sft_data = []
         with open(csv_path, "r") as f:
             reader = csv.DictReader(f)
             for row in reader:
-                data.append({"prompt": row["prompt"], "response": row["response"]})
+                self.sft_data.append({"prompt": row["prompt"], "response": row["response"]})
 
         # Prepare dataset and dataloader
-        dataset = SFTDataset(data, self.tokenizer)
+        dataset = SFTDataset(self.sft_data, self.tokenizer)
         dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)
 
         # Set up optimizer
@@ -188,6 +189,14 @@ class ModelBuilder:
             self.tokenizer.save_pretrained(path)
         st.success("Model saved!")
 
+    def evaluate(self, prompt: str):
+        """Evaluate the model with a prompt"""
+        self.model.eval()
+        with torch.no_grad():
+            inputs = self.tokenizer(prompt, return_tensors="pt").to(self.model.device)
+            outputs = self.model.generate(**inputs, max_new_tokens=50)
+            return self.tokenizer.decode(outputs[0], skip_special_tokens=True)
+
 # Utility Functions
 def sanitize_label(label):
     """Remove invalid characters for Mermaid labels."""
@@ -340,6 +349,7 @@ if st.button("Grow Tree 🌱") and new_node and parent_node:
     # 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.rerun()
 
 # Display Mermaid Diagram
 st.markdown("### Knowledge Tree Visualization")
@@ -429,13 +439,23 @@ if st.button("Predict"):
         with st.expander("Model Configuration", expanded=True):
             base_model = st.selectbox(
                 "Select Base Model",
-                ["distilgpt2", "gpt2", "EleutherAI/pythia-70m"],  # Small models suitable for SFT
+                ["distilgpt2", "gpt2", "EleutherAI/pythia-70m"],
                 help="Choose a small model for fine-tuning"
             )
-            model_name = st.text_input("Model Name", "sft-model")
+            model_name = st.text_input("Model Name", f"sft-model-{int(time.time())}")
             domain = st.text_input("Target Domain", "general")
 
-        # Generate Sample CSV
+        # 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"):
             sample_data = [
                 {"prompt": "What is AI?", "response": "AI is artificial intelligence, simulating human intelligence in machines."},
@@ -449,30 +469,24 @@ if st.button("Predict"):
             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
+        # 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")):
-            config = ModelConfig(
-                name=model_name,
-                base_model=base_model,
-                size="small",
-                domain=domain
-            )
-            builder = ModelBuilder()
+            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)
             
-            # 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)
-                builder.fine_tune_sft(csv_path)
-                builder.save_model(config.model_path)
+                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)
             app_code = f"""
 import streamlit as st
 from transformers import AutoModelForCausalLM, AutoTokenizer
@@ -507,6 +521,23 @@ if st.button("Generate"):
             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__":
-    st.run()
+            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}")