Update app.py

app.py

@@ -2,7 +2,7 @@
 import os
 import shutil
 import glob
-import base64  # Added this import
+import base64
 import streamlit as st
 import pandas as pd
 import torch
@@ -77,13 +77,15 @@ class ModelBuilder:
         self.tokenizer = None
         self.sft_data = None
 
-    def load_model(self, model_path: str):
-        """Load a model from a path"""
+    def load_model(self, model_path: str, config: Optional[ModelConfig] = None):
+        """Load a model from a path with an optional config"""
         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
 
@@ -156,7 +158,9 @@ selected_model = st.sidebar.selectbox("Select Saved Model", ["None"] + model_dir
 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)
+    # Create a config for the loaded model if none exists
+    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()
 
@@ -176,7 +180,7 @@ with tab1:
     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.load_model(base_model, config)  # Pass config here
         builder.save_model(config.model_path)
         st.session_state['builder'] = builder
         st.session_state['model_loaded'] = True
@@ -210,7 +214,12 @@ with tab2:
             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)
+            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)