Update app.py

app.py

@@ -1,5 +1,5 @@
 import streamlit as st
-from transformers import T5ForConditionalGeneration, T5Tokenizer, pipeline
+from transformers import T5ForConditionalGeneration, T5Tokenizer, pipeline, AutoTokenizer, AutoModelForCausalLM
 import torch
 
 # Streamlit app setup
@@ -7,7 +7,12 @@ st.set_page_config(page_title="Chat", layout="wide")
 
 # Sidebar: Model controls
 st.sidebar.title("Model Controls")
-model_name = st.sidebar.text_input("Enter Model Name", value="karthikeyan-r/slm-custom-model_6k")
+model_options = {
+    "1": "karthikeyan-r/slm-custom-model_6k",
+    "2": "karthikeyan-r/calculation_model"
+}
+
+model_choice = st.sidebar.selectbox("Select Model", options=list(model_options.values()))
 load_model_button = st.sidebar.button("Load Model")
 clear_conversation_button = st.sidebar.button("Clear Conversation")
 clear_model_button = st.sidebar.button("Clear Model")
@@ -31,15 +36,39 @@ if "user_input" not in st.session_state:
 if load_model_button:
     with st.spinner("Loading model..."):
         try:
-            device = 0 if torch.cuda.is_available() else -1
-            st.session_state["model"] = T5ForConditionalGeneration.from_pretrained(model_name, cache_dir="./model_cache")
-            st.session_state["tokenizer"] = T5Tokenizer.from_pretrained(model_name, cache_dir="./model_cache")
-            st.session_state["qa_pipeline"] = pipeline(
-                "text2text-generation",
-                model=st.session_state["model"],
-                tokenizer=st.session_state["tokenizer"],
-                device=device
-            )
+            # 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")
+                
+                # Add special tokens if not present
+                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))
+
+                # Update configuration
+                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  # Calculation model doesn't use text2text pipeline
             st.success("Model loaded successfully and ready!")
         except Exception as e:
             st.error(f"Error loading model: {e}")
@@ -74,7 +103,7 @@ if st.session_state["qa_pipeline"]:
         if user_input:
             with st.spinner("Generating response..."):
                 try:
-                    # Generate the model response
+                    # Generate the model response for general QA (T5 model)
                     response = st.session_state["qa_pipeline"](f"Q: {user_input}", max_length=400)
                     generated_text = response[0]["generated_text"]
 
@@ -89,6 +118,49 @@ if st.session_state["qa_pipeline"]:
                     display_conversation()
                 except Exception as e:
                     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"]:
+        user_input = st.text_input(
+            "Enter your query for calculation:",
+            value=st.session_state["user_input"],
+            key="calculation_input",
+        )
+        if st.button("Send Calculation", key="send_calculation_button"):
+            if user_input:
+                with st.spinner("Generating response..."):
+                    try:
+                        # Generate the model response for the calculation model
+                        inputs = st.session_state["tokenizer"](f"Input: {user_input}\nOutput:", return_tensors="pt", padding=True, truncation=True)
+                        input_ids = inputs.input_ids
+                        attention_mask = inputs.attention_mask
+
+                        output = st.session_state["model"].generate(
+                            input_ids=input_ids,
+                            attention_mask=attention_mask,
+                            max_length=50,
+                            pad_token_id=st.session_state["tokenizer"].pad_token_id,
+                            eos_token_id=st.session_state["tokenizer"].eos_token_id,
+                            do_sample=False
+                        )
+
+                        decoded_output = st.session_state["tokenizer"].decode(output[0], skip_special_tokens=True)
+                        if "Output:" in decoded_output:
+                            answer = decoded_output.split("Output:")[-1].strip()
+                        else:
+                            answer = decoded_output.strip()
+
+                        # Update the conversation
+                        st.session_state["conversation"].append(("You", user_input))
+                        st.session_state["conversation"].append(("Model", answer))
+
+                        # Clear the input field after submission
+                        st.session_state["user_input"] = ""
+
+                        # Rerender the conversation immediately
+                        display_conversation()
+                    except Exception as e:
+                        st.error(f"Error generating response: {e}")
 
 # Clear Conversation
 if clear_conversation_button: