Update app.py

app.py

@@ -1,26 +1,32 @@
 import streamlit as st
-from transformers import T5ForConditionalGeneration, T5Tokenizer, pipeline, AutoTokenizer, AutoModelForCausalLM
+from transformers import (
+    T5ForConditionalGeneration,
+    T5Tokenizer,
+    pipeline,
+    AutoTokenizer,
+    AutoModelForCausalLM
+)
 import torch
 
-# Streamlit app setup
+# ----- Streamlit page config -----
 st.set_page_config(page_title="Chat", layout="wide")
 
-# Sidebar: Model controls
+# ----- Sidebar: Model controls -----
 st.sidebar.title("Model Controls")
 model_options = {
     "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()))
+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")
 
-# Main UI
-st.title("Chat Conversation UI")
-
-# Session states
+# ----- Session States -----
 if "model" not in st.session_state:
     st.session_state["model"] = None
 if "tokenizer" not in st.session_state:
@@ -28,144 +34,148 @@ 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": "user"/"assistant", "content": "..."}]
     st.session_state["conversation"] = []
-if "user_input" not in st.session_state:
-    st.session_state["user_input"] = ""
 
-# Load Model
+# ----- Load Model -----
 if load_model_button:
     with st.spinner("Loading model..."):
         try:
-            # Load the selected model
             if model_choice == model_options["1"]:
-                
-                # Load the calculation model (calculation_model)
+                # 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 not present
+
+                # 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))
 
-                # 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.session_state["qa_pipeline"] = None  # Not needed for calculation model
+
             elif model_choice == model_options["2"]:
-                # Load the T5 model for general QA (slm-custom-model_6k)
+                # Load the T5 model for general QA
                 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(
+                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=st.session_state["model"],
-                    tokenizer=st.session_state["tokenizer"],
+                    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}")
 
-# Clear Model
+# ----- Clear Model -----
 if clear_model_button:
     st.session_state["model"] = None
     st.session_state["tokenizer"] = None
     st.session_state["qa_pipeline"] = None
     st.success("Model cleared.")
 
-# Chat Conversation Display
-def display_conversation():
-    """Display the chat conversation dynamically."""
-    st.subheader("Conversation")
-    for idx, (speaker, message) in enumerate(st.session_state["conversation"]):
-        if speaker == "You":
-            st.markdown(f"**You:** {message}")
-        else:
-            st.markdown(f"**Model:** {message}")
+# ----- Clear Conversation -----
+if clear_conversation_button:
+    st.session_state["conversation"] = []
+    st.success("Conversation cleared.")
 
-display_conversation()
+# ----- Display Chat Conversation -----
+st.title("Chat Conversation UI")
 
-# Input Area
+# Loop through existing conversation in session_state and display it
+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"])
+
+# ----- Chat Input Logic -----
+# If we have a T5 pipeline (general QA model):
 if st.session_state["qa_pipeline"]:
-    user_input = st.text_input(
-        "Enter your query:",
-        value=st.session_state["user_input"],  # Use session state for persistence
-        key="chat_input",
-    )
-    if st.button("Send", key="send_button"):
-        if user_input:
+    # Use the new Streamlit chat input
+    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 response
+        with st.chat_message("assistant"):
             with st.spinner("Generating response..."):
                 try:
-                    # Generate the model response for general QA (T5 model)
                     response = st.session_state["qa_pipeline"](f"Q: {user_input}", max_length=250)
                     generated_text = response[0]["generated_text"]
+                except Exception as e:
+                    generated_text = f"Error: {str(e)}"
+
+            st.write(generated_text)
 
-                    # Update the conversation
-                    st.session_state["conversation"].append(("You", user_input))
-                    st.session_state["conversation"].append(("Model", generated_text))
+        # 3) Save assistant message
+        st.session_state["conversation"].append({"role": "assistant", "content": generated_text})
 
-                    # Clear the input field after submission
-                    st.session_state["user_input"] = ""
+# If we have the calculation model loaded (model_options["1"]):
+elif st.session_state["model"] and (model_choice == model_options["1"]):
+    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 response
+        with st.chat_message("assistant"):
+            with st.spinner("Generating response..."):
+                try:
+                    tokenizer = st.session_state["tokenizer"]
+                    model = st.session_state["model"]
+
+                    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()
 
-                    # Rerender the conversation immediately
-                    display_conversation()
                 except Exception as e:
-                    st.error(f"Error generating response: {e}")
+                    answer = f"Error: {str(e)}"
+
+            st.write(answer)
+
+        # 3) Save assistant message
+        st.session_state["conversation"].append({"role": "assistant", "content": answer})
 else:
-    # Handle user input for the calculation model (calculation_model)
-    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"],
-            key="calculation_input",
-        )
-        if st.button("Submit", 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=250,
-                            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:
-    st.session_state["conversation"] = []
-    st.session_state["user_input"] = ""  # Clear input field
-    st.success("Conversation cleared.")
+    # If no model is loaded at all
+    st.info("No model is loaded. Please select a model and click 'Load Model' from the sidebar.")