Update app.py

app.py

@@ -14,13 +14,20 @@ from dataclasses import dataclass
 from typing import Optional, Tuple
 import zipfile
 import math
+from PIL import Image
+import random
 
-# Page Configuration
+# Page Configuration with a Dash of Humor
 st.set_page_config(
-    page_title="SFT Model Builder 🚀",
+    page_title="SFT Tiny Titans 🚀",
     page_icon="🤖",
     layout="wide",
     initial_sidebar_state="expanded",
+    menu_items={
+        'Get Help': 'https://huggingface.co/awacke1',
+        'Report a bug': 'https://huggingface.co/spaces/awacke1',
+        'About': "Tiny Titans: Small models, big dreams, and a sprinkle of chaos! 🌌"
+    }
 )
 
 # Model Configuration Class
@@ -49,53 +56,37 @@ class SFTDataset(Dataset):
         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"
-        )
-        
+        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"
-        )
+        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
+        labels[:prompt_len] = -100
         
-        return {
-            "input_ids": input_ids,
-            "attention_mask": attention_mask,
-            "labels": labels
-        }
+        return {"input_ids": input_ids, "attention_mask": attention_mask, "labels": labels}
 
-# Model Builder Class
+# Model Builder Class with Easter Egg Jokes
 class ModelBuilder:
     def __init__(self):
         self.config = None
         self.model = None
         self.tokenizer = None
         self.sft_data = None
+        self.jokes = ["Why did the AI go to therapy? Too many layers to unpack! 😂", "Training complete! Time for a binary coffee break. ☕"]
 
     def load_model(self, model_path: str, config: Optional[ModelConfig] = None):
-        with st.spinner("Loading model... ⏳"):
+        with st.spinner(f"Loading {model_path}... ⏳ (Patience, young padawan!)"):
             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! ✅")
+        st.success(f"Model loaded! 🎉 {random.choice(self.jokes)}")
         return self
 
     def fine_tune_sft(self, csv_path: str, epochs: int = 3, batch_size: int = 4):
@@ -113,7 +104,7 @@ class ModelBuilder:
         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}... ⚙️"):
+            with st.spinner(f"Training epoch {epoch + 1}/{epochs}... ⚙️ (The AI is lifting weights!)"):
                 total_loss = 0
                 for batch in dataloader:
                     optimizer.zero_grad()
@@ -126,35 +117,29 @@ class ModelBuilder:
                     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(f"SFT Fine-tuning completed! 🎉 {random.choice(self.jokes)}")
         return self
 
     def save_model(self, path: str):
-        with st.spinner("Saving model... 💾"):
+        with st.spinner("Saving model... 💾 (Packing the AI’s suitcase!)"):
             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}! ✅")
+        st.success(f"Model saved at {path}! ✅ May the force be with it.")
 
     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
-            )
+            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
+# Utility Functions with Wit
 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>'
+    return f'<a href="data:{mime_type};base64,{b64}" download="{os.path.basename(file_path)}">{label} 📥 (Grab it before it runs away!)</a>'
 
 def zip_directory(directory_path, zip_path):
     with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf:
@@ -167,12 +152,14 @@ def zip_directory(directory_path, zip_path):
 def get_model_files():
     return [d for d in glob.glob("models/*") if os.path.isdir(d)]
 
+def get_gallery_files(file_types):
+    files = []
+    for ext in file_types:
+        files.extend(glob.glob(f"*.{ext}"))
+    return sorted(files)
+
 # 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 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)
@@ -188,13 +175,28 @@ def calculate_cargo_travel_time(
     return round(flight_time, 2)
 
 # Main App
-st.title("SFT Model Builder 🤖🚀")
-
-# Sidebar for Model Management
-st.sidebar.header("Model Management 🗂️")
+st.title("SFT Tiny Titans 🚀 (Small but Mighty!)")
+
+# Sidebar with Galleries
+st.sidebar.header("Galleries & Shenanigans 🎨")
+st.sidebar.subheader("Image Gallery 📸")
+img_files = get_gallery_files(["png", "jpg", "jpeg"])
+if img_files:
+    img_cols = st.sidebar.slider("Image Columns 📸", 1, 5, 3)
+    cols = st.sidebar.columns(img_cols)
+    for idx, img_file in enumerate(img_files[:img_cols * 2]):  # Limit to 2 rows
+        with cols[idx % img_cols]:
+            st.image(Image.open(img_file), caption=f"{img_file} 🖼", use_column_width=True)
+
+st.sidebar.subheader("CSV Gallery 📊")
+csv_files = get_gallery_files(["csv"])
+if csv_files:
+    for csv_file in csv_files[:5]:  # Limit to 5
+        st.sidebar.markdown(get_download_link(csv_file, "text/csv", f"{csv_file} 📊"), unsafe_allow_html=True)
+
+st.sidebar.subheader("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()
@@ -204,21 +206,16 @@ if selected_model != "None" and st.sidebar.button("Load Model 📂"):
     st.rerun()
 
 # Main UI with Tabs
-tab1, tab2, tab3, tab4 = st.tabs(["Build New Model 🌱", "Fine-Tune Model 🔧", "Test Model 🧪", "Agentic RAG Demo 🌐"])
+tab1, tab2, tab3, tab4 = st.tabs(["Build Tiny Titan 🌱", "Fine-Tune Titan 🔧", "Test Titan 🧪", "Agentic RAG Party 🌐"])
 
 with tab1:
-    st.header("Build New Model 🌱")
+    st.header("Build Tiny Titan 🌱 (Assemble Your Mini-Mecha!)")
     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)"
+        "Select Tiny Model",
+        ["HuggingFaceTB/SmolLM-135M", "HuggingFaceTB/SmolLM-360M", "Qwen/Qwen1.5-0.5B-Chat"],
+        help="Pick a pint-sized powerhouse (<1 GB)! SmolLM-135M (~270 MB), SmolLM-360M (~720 MB), Qwen1.5-0.5B (~1 GB)"
     )
-    model_name = st.text_input("Model Name", f"new-model-{int(time.time())}")
+    model_name = st.text_input("Model Name", f"tiny-titan-{int(time.time())}")
     domain = st.text_input("Target Domain", "general")
 
     if st.button("Download Model ⬇️"):
@@ -228,19 +225,19 @@ with tab1:
         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.success(f"Model downloaded and saved to {config.model_path}! 🎉 (Tiny but feisty!)")
         st.rerun()
 
 with tab2:
-    st.header("Fine-Tune Model 🔧")
+    st.header("Fine-Tune Titan 🔧 (Teach Your Titan Some Tricks!)")
     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! ⚠️")
+        st.warning("Please build or load a Titan first! ⚠️ (No Titan, no party!)")
     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."},
+                {"prompt": "What is AI?", "response": "AI is artificial intelligence, simulating human smarts in machines."},
+                {"prompt": "Explain machine learning", "response": "Machine learning is AI’s gym where models bulk up on data."},
+                {"prompt": "What is a neural network?", "response": "A neural network is a brainy AI mimicking human noggins."},
             ]
             csv_path = f"sft_data_{int(time.time())}.csv"
             with open(csv_path, "w", newline="") as f:
@@ -248,7 +245,7 @@ with tab2:
                 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}! ✅")
+            st.success(f"Sample CSV generated as {csv_path}! ✅ (Fresh from the data oven!)")
 
         uploaded_csv = st.file_uploader("Upload CSV for SFT", type="csv")
         if uploaded_csv and st.button("Fine-Tune with Uploaded CSV 🔄"):
@@ -263,20 +260,20 @@ with tab2:
                 domain=st.session_state['builder'].config.domain
             )
             st.session_state['builder'].config = new_config
-            with st.status("Fine-tuning model... ⏳", expanded=True) as status:
+            with st.status("Fine-tuning Titan... ⏳ (Whipping it into shape!)", 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")
+                status.update(label="Fine-tuning completed! 🎉 (Titan’s ready to rumble!)", 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.markdown(get_download_link(zip_path, "application/zip", "Download Fine-Tuned Titan"), unsafe_allow_html=True)
             st.rerun()
 
 with tab3:
-    st.header("Test Model 🧪")
+    st.header("Test Titan 🧪 (Put Your Titan to the Test!)")
     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! ⚠️")
+        st.warning("Please build or load a Titan first! ⚠️ (No Titan, no test drive!)")
     else:
         if st.session_state['builder'].sft_data:
             st.write("Testing with SFT Data:")
@@ -286,15 +283,15 @@ with tab3:
                 generated = st.session_state['builder'].evaluate(prompt)
                 st.write(f"**Prompt**: {prompt}")
                 st.write(f"**Expected**: {expected}")
-                st.write(f"**Generated**: {generated}")
+                st.write(f"**Generated**: {generated} (Titan says: '{random.choice(['Bleep bloop!', 'I am groot!', '42!'])}')")
                 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}")
+            st.write(f"**Generated Response**: {result} (Titan’s wisdom unleashed!)")
 
-        if st.button("Export Model Files 📦"):
+        if st.button("Export Titan Files 📦"):
             config = st.session_state['builder'].config
             app_code = f"""
 import streamlit as st
@@ -303,47 +300,47 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
 model = AutoModelForCausalLM.from_pretrained("{config.model_path}")
 tokenizer = AutoTokenizer.from_pretrained("{config.model_path}")
 
-st.title("SFT Model Demo")
+st.title("Tiny Titan 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:
+            with open("titan_app.py", "w") as f:
                 f.write(app_code)
             reqs = "streamlit\ntorch\ntransformers\n"
-            with open("sft_requirements.txt", "w") as f:
+            with open("titan_requirements.txt", "w") as f:
                 f.write(reqs)
             readme = f"""
-# SFT Model Demo
+# Tiny Titan 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".
+1. Install requirements: `pip install -r titan_requirements.txt`
+2. Run the app: `streamlit run titan_app.py`
+3. Input a prompt and click "Generate". Watch the magic unfold! 🪄
 """
-            with open("sft_README.md", "w") as f:
+            with open("titan_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! ✅")
+            st.markdown(get_download_link("titan_app.py", "text/plain", "Download App"), unsafe_allow_html=True)
+            st.markdown(get_download_link("titan_requirements.txt", "text/plain", "Download Requirements"), unsafe_allow_html=True)
+            st.markdown(get_download_link("titan_README.md", "text/markdown", "Download README"), unsafe_allow_html=True)
+            st.success("Titan files exported! ✅ (Ready to conquer the galaxy!)")
 
 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.")
+    st.header("Agentic RAG Party 🌐 (Party Like It’s 2099!)")
+    st.write("This demo uses tiny Titans with Agentic RAG to plan a superhero party, powered by DuckDuckGo retrieval!")
 
     if st.button("Run Agentic RAG Demo 🎉"):
         try:
             from smolagents import CodeAgent, DuckDuckGoSearchTool, VisitWebpageTool
 
-            # Load selected tiny model
+            # Load a tiny model (default to SmolLM-135M for speed)
             tokenizer = AutoTokenizer.from_pretrained("HuggingFaceTB/SmolLM-135M")
             model = AutoModelForCausalLM.from_pretrained("HuggingFaceTB/SmolLM-135M")
 
-            # Define Agentic RAG agent
+            # Define Agentic RAG agent with a witty twist
             agent = CodeAgent(
                 model=model,
                 tokenizer=tokenizer,
@@ -355,16 +352,18 @@ with tab4:
             )
             
             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.
+Plan a luxury superhero-themed party at Wayne Manor (42.3601° N, 71.0589° W). Use DuckDuckGo to search for the latest superhero party trends,
+refine results for luxury elements (decorations, entertainment, catering), and calculate cargo travel times from key locations 
+(New York: 40.7128° N, 74.0060° W; LA: 34.0522° N, 118.2437° W; London: 51.5074° N, 0.1278° W) to Wayne Manor. 
+Synthesize a plan with at least 6 entries in a pandas dataframe, including locations, travel times, and luxury ideas.
+Add a random superhero catchphrase to each entry for fun!
 """
-            with st.spinner("Running Agentic RAG system... ⏳"):
+            with st.spinner("Planning the ultimate superhero bash... ⏳ (Calling all caped crusaders!)"):
                 result = agent.run(task)
-            st.write("Agentic RAG Result:")
+            st.write("Agentic RAG Party Plan:")
             st.write(result)
+            st.write("Party on, Wayne! 🦸‍♂️🎉")
         except ImportError:
             st.error("Please install required packages: `pip install smolagents pandas`")
         except Exception as e:
-            st.error(f"Error running demo: {str(e)}")
+            st.error(f"Error running demo: {str(e)} (Even Batman has off days!)")