Update app.py

app.py

@@ -2,7 +2,7 @@ import os
 import nest_asyncio
 nest_asyncio.apply()
 import streamlit as st
-from transformers import pipeline
+from transformers import pipeline, AutoTokenizer
 from huggingface_hub import login
 from streamlit.components.v1 import html
 import pandas as pd
@@ -53,7 +53,7 @@ st.set_page_config(page_title="Review Scorer & Report Generator", page_icon="
 st.header("Review Scorer & Report Generator")
 
 # Concise introduction
-st.write("This model will score your reviews in your CSV file and generate a report based on those results.")
+st.write("This model will score your reviews in your CSV file and generate a report based on your query and those results.")
 
 # Load models with caching to avoid reloading on every run
 @st.cache_resource
@@ -65,15 +65,18 @@ def load_models():
         st.info("Loading sentiment analysis model...")
         score_pipe = pipeline("text-classification", 
                               model="nlptown/bert-base-multilingual-uncased-sentiment", 
-                              device=0)
+                              device=0 if torch.cuda.is_available() else -1)
         st.success("Sentiment analysis model loaded successfully!")
     except Exception as e:
         st.error(f"Error loading score model: {e}")
     
     try:
         st.info("Loading Gemma model...")
+        # Load the tokenizer separately with the chat template
+        tokenizer = AutoTokenizer.from_pretrained("google/gemma-3-1b-it")
         gemma_pipe = pipeline("text-generation", 
-                              model="google/gemma-2-2b-it",
+                              model="google/gemma-3-1b-it",
+                              tokenizer=tokenizer,  # Pass the loaded tokenizer here
                               device=0, 
                               torch_dtype=torch.bfloat16)
         st.success("Gemma model loaded successfully!")
@@ -83,9 +86,34 @@ def load_models():
     
     return score_pipe, gemma_pipe
 
+def extract_assistant_content(raw_response):
+    """Extract only the assistant's content from the Gemma-3 response."""
+    # Convert to string and work with it directly
+    response_str = str(raw_response)
+
+    # Look for the assistant's content marker
+    assistant_marker = "'role': 'assistant', 'content': '"
+    if assistant_marker in response_str:
+        start_idx = response_str.find(assistant_marker) + len(assistant_marker)
+        # Extract everything after the marker until the end or closing quote
+        content = response_str[start_idx:]
+
+        # Find the end of the content (last single quote before the end of the string or before closing curly brace)
+        end_markers = ["'}", "'}]"]
+        end_idx = len(content)
+        for marker in end_markers:
+            pos = content.rfind(marker)
+            if pos != -1 and pos < end_idx:
+                end_idx = pos
+
+        return content[:end_idx]
+
+    # Fallback - return the original response
+    return response_str
 
 score_pipe, gemma_pipe = load_models()
 
+
 # Input: Query text for scoring and CSV file upload for candidate reviews
 query_input = st.text_area("Enter your query text for analysis (this does not need to be part of the CSV):")
 uploaded_file = st.file_uploader("Upload Reviews CSV File (must contain a 'reviewText' column)", type=["csv"])
@@ -148,10 +176,11 @@ Candidate Reviews with their scores:
         """}
             ]
             
-            output = gemma_pipe(messages, max_new_tokens=50)
+            raw_result = gemma_pipe(messages, max_new_tokens=50)
+            report = extract_assistant_content(raw_result)
             progress_bar.progress(100)
             status_text.success("**✅ Generation complete!**")
             html("<script>localStorage.setItem('freezeTimer', 'true');</script>", height=0)
             st.session_state.timer_frozen = True
             #st.write("**Scored Candidate Reviews:**", scored_docs)
-            st.write("**Generated Report:**", output)
+            st.write("**Generated Report:**", report)