Update app.py

app.py

@@ -34,59 +34,54 @@ if "tokenizer" not in st.session_state:
 if "qa_pipeline" not in st.session_state:
     st.session_state["qa_pipeline"] = None
 if "conversation" not in st.session_state:
-    # We'll store conversation as a list of dicts,
-    # e.g. [{"role": "assistant", "content": "Hello..."}, {"role": "user", "content": "..."}]
     st.session_state["conversation"] = []
 
 # ----- Load Model -----
+def load_model():
+    if st.session_state["model"] is None or st.session_state["tokenizer"] is None:
+        with st.spinner("Loading model..."):
+            try:
+                if model_choice == model_options["1"]:
+                    # Load the calculation model
+                    tokenizer = AutoTokenizer.from_pretrained(model_choice, cache_dir="./model_cache")
+                    model = AutoModelForCausalLM.from_pretrained(model_choice, cache_dir="./model_cache")
+
+                    # Add special tokens if needed
+                    if tokenizer.pad_token is None:
+                        tokenizer.add_special_tokens({'pad_token': '[PAD]'})
+                        model.resize_token_embeddings(len(tokenizer))
+                    if tokenizer.eos_token is None:
+                        tokenizer.add_special_tokens({'eos_token': '[EOS]'})
+                        model.resize_token_embeddings(len(tokenizer))
+
+                    model.config.pad_token_id = tokenizer.pad_token_id
+                    model.config.eos_token_id = tokenizer.eos_token_id
+
+                    st.session_state["model"] = model
+                    st.session_state["tokenizer"] = tokenizer
+                    st.session_state["qa_pipeline"] = None  # Not needed for calculation model
+
+                elif model_choice == model_options["2"]:
+                    # Load the T5 model for general QA
+                    device = 0 if torch.cuda.is_available() else -1
+                    model = T5ForConditionalGeneration.from_pretrained(model_choice, cache_dir="./model_cache")
+                    tokenizer = T5Tokenizer.from_pretrained(model_choice, cache_dir="./model_cache")
+                    qa_pipe = pipeline(
+                        "text2text-generation",
+                        model=model,
+                        tokenizer=tokenizer,
+                        device=device
+                    )
+                    st.session_state["model"] = model
+                    st.session_state["tokenizer"] = tokenizer
+                    st.session_state["qa_pipeline"] = qa_pipe
+
+                st.success("Model loaded successfully and ready!")
+            except Exception as e:
+                st.error(f"Error loading model: {e}")
+
 if load_model_button:
-    with st.spinner("Loading model..."):
-        try:
-            if model_choice == model_options["1"]:
-                # Load the calculation model
-                tokenizer = AutoTokenizer.from_pretrained(model_choice, cache_dir="./model_cache")
-                model = AutoModelForCausalLM.from_pretrained(model_choice, cache_dir="./model_cache")
-
-                # Add special tokens if needed
-                if tokenizer.pad_token is None:
-                    tokenizer.add_special_tokens({'pad_token': '[PAD]'})
-                    model.resize_token_embeddings(len(tokenizer))
-                if tokenizer.eos_token is None:
-                    tokenizer.add_special_tokens({'eos_token': '[EOS]'})
-                    model.resize_token_embeddings(len(tokenizer))
-
-                model.config.pad_token_id = tokenizer.pad_token_id
-                model.config.eos_token_id = tokenizer.eos_token_id
-
-                st.session_state["model"] = model
-                st.session_state["tokenizer"] = tokenizer
-                st.session_state["qa_pipeline"] = None  # Not needed for calculation model
-
-            elif model_choice == model_options["2"]:
-                # Load the T5 model for general QA
-                device = 0 if torch.cuda.is_available() else -1
-                model = T5ForConditionalGeneration.from_pretrained(model_choice, cache_dir="./model_cache")
-                tokenizer = T5Tokenizer.from_pretrained(model_choice, cache_dir="./model_cache")
-                qa_pipe = pipeline(
-                    "text2text-generation",
-                    model=model,
-                    tokenizer=tokenizer,
-                    device=device
-                )
-                st.session_state["model"] = model
-                st.session_state["tokenizer"] = tokenizer
-                st.session_state["qa_pipeline"] = qa_pipe
-
-            # If conversation is empty, insert a welcome message
-            if len(st.session_state["conversation"]) == 0:
-                st.session_state["conversation"].append({
-                    "role": "assistant",
-                    "content": "Hello! I’m your assistant. How can I help you today?"
-                })
-
-            st.success("Model loaded successfully and ready!")
-        except Exception as e:
-            st.error(f"Error loading model: {e}")
+    load_model()
 
 # ----- Clear Model -----
 if clear_model_button:
@@ -103,93 +98,42 @@ if clear_conversation_button:
 # ----- Title -----
 st.title("Chat Conversation UI")
 
-
-user_input = None
-
-if st.session_state["qa_pipeline"]:
-    # T5 pipeline
-    user_input = st.chat_input("Enter your query:")
-    if user_input:
-        # 1) Save user message
-        st.session_state["conversation"].append({
-            "role": "user",
-            "content": user_input
-        })
-
-        # 2) Generate assistant response
+# ----- User Input and Processing -----
+user_input = st.chat_input("Enter your query:")
+if user_input:
+    # Save user input
+    st.session_state["conversation"].append({
+        "role": "user",
+        "content": user_input
+    })
+
+    # Generate response
+    if st.session_state["qa_pipeline"]:
         try:
-            response = st.session_state["qa_pipeline"](
-                f"Q: {user_input}", max_length=250
-            )
+            response = st.session_state["qa_pipeline"](f"Q: {user_input}", max_length=250)
             answer = response[0]["generated_text"]
         except Exception as e:
             answer = f"Error: {str(e)}"
-
-        # 3) Append assistant message to conversation
-        st.session_state["conversation"].append({
-            "role": "assistant",
-            "content": answer
-        })
-
-elif st.session_state["model"] and (model_choice == model_options["1"]):
-    # Calculation model
-    user_input = st.chat_input("Enter your query for calculation:")
-    if user_input:
-        # 1) Save user message
-        st.session_state["conversation"].append({
-            "role": "user",
-            "content": user_input
-        })
-
-        # 2) Generate assistant response
-        tokenizer = st.session_state["tokenizer"]
-        model = st.session_state["model"]
-
+    elif st.session_state["model"] and model_choice == model_options["1"]:
         try:
-            inputs = tokenizer(
-                f"Input: {user_input}\nOutput:",
-                return_tensors="pt",
-                padding=True,
-                truncation=True
-            )
-            input_ids = inputs.input_ids
-            attention_mask = inputs.attention_mask
-
-            output = model.generate(
-                input_ids=input_ids,
-                attention_mask=attention_mask,
-                max_length=250,
-                pad_token_id=tokenizer.pad_token_id,
-                eos_token_id=tokenizer.eos_token_id,
-                do_sample=False
-            )
-
-            decoded_output = tokenizer.decode(
-                output[0],
-                skip_special_tokens=True
-            )
-            # Extract answer after 'Output:' if present
-            if "Output:" in decoded_output:
-                answer = decoded_output.split("Output:")[-1].strip()
-            else:
-                answer = decoded_output.strip()
+            tokenizer = st.session_state["tokenizer"]
+            model = st.session_state["model"]
+
+            inputs = tokenizer(f"Input: {user_input}\nOutput:", return_tensors="pt", padding=True, truncation=True)
+            output = model.generate(inputs.input_ids, max_length=250, pad_token_id=tokenizer.pad_token_id)
+            answer = tokenizer.decode(output[0], skip_special_tokens=True).split("Output:")[-1].strip()
         except Exception as e:
             answer = f"Error: {str(e)}"
+    else:
+        answer = "No model is loaded. Please select and load a model."
 
-        # 3) Append assistant message to conversation
-        st.session_state["conversation"].append({
-            "role": "assistant",
-            "content": answer
-        })
-else:
-    # If no model is loaded:
-    st.info("No model is loaded. Please select a model and click 'Load Model' from the sidebar.")
-
+    # Save assistant response
+    st.session_state["conversation"].append({
+        "role": "assistant",
+        "content": answer
+    })
 
+# Display conversation
 for message in st.session_state["conversation"]:
-    if message["role"] == "user":
-        with st.chat_message("user"):
-            st.write(message["content"])
-    else:
-        with st.chat_message("assistant"):
-            st.write(message["content"])
+    with st.chat_message(message["role"]):
+        st.write(message["content"])