Update app.py

app.py

@@ -8,8 +8,8 @@ st.set_page_config(page_title="Chat", layout="wide")
 # Sidebar: Model controls
 st.sidebar.title("Model Controls")
 model_options = {
-    "1": "karthikeyan-r/slm-custom-model_6k",
-    "2": "karthikeyan-r/calculation_model"
+    "1": "karthikeyan-r/calculation_model_11k",
+    "2": "karthikeyan-r/slm-custom-model_6k"
 }
 
 model_choice = st.sidebar.selectbox("Select Model", options=list(model_options.values()))
@@ -38,17 +38,7 @@ if load_model_button:
         try:
             # Load the selected model
             if model_choice == model_options["1"]:
-                # Load the T5 model for general QA (slm-custom-model_6k)
-                device = 0 if torch.cuda.is_available() else -1
-                st.session_state["model"] = T5ForConditionalGeneration.from_pretrained(model_choice, cache_dir="./model_cache")
-                st.session_state["tokenizer"] = T5Tokenizer.from_pretrained(model_choice, cache_dir="./model_cache")
-                st.session_state["qa_pipeline"] = pipeline(
-                    "text2text-generation",
-                    model=st.session_state["model"],
-                    tokenizer=st.session_state["tokenizer"],
-                    device=device
-                )
-            elif model_choice == model_options["2"]:
+                
                 # Load the calculation model (calculation_model)
                 tokenizer = AutoTokenizer.from_pretrained(model_choice, cache_dir="./model_cache")
                 model = AutoModelForCausalLM.from_pretrained(model_choice, cache_dir="./model_cache")
@@ -69,6 +59,18 @@ if load_model_button:
                 st.session_state["model"] = model
                 st.session_state["tokenizer"] = tokenizer
                 st.session_state["qa_pipeline"] = None  # Calculation model doesn't use text2text pipeline
+                
+            elif model_choice == model_options["2"]:
+                # Load the T5 model for general QA (slm-custom-model_6k)
+                device = 0 if torch.cuda.is_available() else -1
+                st.session_state["model"] = T5ForConditionalGeneration.from_pretrained(model_choice, cache_dir="./model_cache")
+                st.session_state["tokenizer"] = T5Tokenizer.from_pretrained(model_choice, cache_dir="./model_cache")
+                st.session_state["qa_pipeline"] = pipeline(
+                    "text2text-generation",
+                    model=st.session_state["model"],
+                    tokenizer=st.session_state["tokenizer"],
+                    device=device
+                )
             st.success("Model loaded successfully and ready!")
         except Exception as e:
             st.error(f"Error loading model: {e}")
@@ -120,7 +122,7 @@ if st.session_state["qa_pipeline"]:
                     st.error(f"Error generating response: {e}")
 else:
     # Handle user input for the calculation model (calculation_model)
-    if st.session_state["model"] and model_choice == model_options["2"]:
+    if st.session_state["model"] and model_choice == model_options["1"]:
         user_input = st.text_input(
             "Enter your query for calculation:",
             value=st.session_state["user_input"],