Update src/app.py

lowering the first token

src/app.py

@@ -1,189 +1,200 @@
-    import streamlit as st
-    import torch
-    from transformers import AutoModelForCausalLM, AutoTokenizer
-    import numpy as np
-    import colorsys
-    import math
-    
-    # --- Core Functions (Cached) ---
-    @st.cache_resource
-    def load_model(model_name):
-        """Loads the specified model and tokenizer from Hugging Face."""
-        import os
-        
-        # Set cache directory to writable location
-        os.environ['HF_HOME'] = '/tmp/hf_cache'
-        os.environ['TRANSFORMERS_CACHE'] = '/tmp/hf_cache'
-        
-        st.info(f"Loading model '{model_name}'... This may take a moment on first run.")
-        
-        try:
-            tokenizer = AutoTokenizer.from_pretrained(
-                model_name,
-                cache_dir="/tmp/hf_cache",
-                trust_remote_code=True
-            )
-            model = AutoModelForCausalLM.from_pretrained(
-                model_name,
-                cache_dir="/tmp/hf_cache",
-                trust_remote_code=True
-            )
-            st.success(f"Model '{model_name}' loaded and ready!")
-            return tokenizer, model
-            
-        except PermissionError as e:
-            st.error(f"Permission error: {e}")
-            st.error("Try refreshing the page or waiting a moment if another download is in progress.")
-            st.stop()
-            
-        except Exception as e:
-            st.error(f"An error occurred loading the model: {e}")
-            st.stop()
-    
-    # --- Analysis and Helper Functions ---
-    def get_analysis_data(text_to_analyze, system_prompt, tokenizer, model):
-        """Calculates log-likelihood and confidence for each token."""
-        messages = [
-            {"role": "system", "content": system_prompt},
-            {"role": "user", "content": text_to_analyze},
-        ]
-        tokenized_chat = tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=False, return_tensors="pt")
-        user_text_token_ids = tokenizer.encode(text_to_analyze, add_special_tokens=False)
-        full_ids_list = tokenized_chat[0].tolist()
-        start_index = -1
-        for i in range(len(full_ids_list) - len(user_text_token_ids) + 1):
-            if full_ids_list[i:i+len(user_text_token_ids)] == user_text_token_ids:
-                start_index = i
-                break
-        if start_index == -1: return []
-        end_index = start_index + len(user_text_token_ids)
-    
-        with torch.no_grad():
-            outputs = model(tokenized_chat, labels=tokenized_chat)
-            logits = outputs.logits
-        log_probs = torch.nn.functional.log_softmax(logits, dim=-1)
-        probs = torch.nn.functional.softmax(logits, dim=-1)
-        sliced_log_probs = log_probs[0, start_index-1:end_index-1, :]
-        sliced_probs = probs[0, start_index-1:end_index-1, :]
-        sliced_tokens = tokenized_chat[0, start_index:end_index]
-        sequence_log_probs = sliced_log_probs.gather(1, sliced_tokens.unsqueeze(-1)).squeeze().tolist()
-        sequence_probs = sliced_probs.gather(1, sliced_tokens.unsqueeze(-1)).squeeze().tolist()
-        tokens = tokenizer.convert_ids_to_tokens(sliced_tokens)
-        if not isinstance(sequence_log_probs, list):
-            sequence_log_probs, sequence_probs = [sequence_log_probs], [sequence_probs]
-        return list(zip(tokens, sequence_log_probs, sequence_probs))
-    
-    def find_high_perplexity_phrases(analysis_data, std_dev_threshold=1.5):
-        """Identifies and groups high-perplexity tokens into phrases."""
-        if not analysis_data: return []
-        log_probs = [lp for _, lp, _ in analysis_data]
-        mean_lp = np.mean(log_probs)
-        std_lp = np.std(log_probs)
-        threshold = mean_lp - std_dev_threshold * std_lp
-        
-        outlier_phrases = []
-        current_phrase = ""
-        for token, log_prob, _ in analysis_data:
-            display_token = token.replace('Ġ', ' ').replace(' ', ' ')
-            if log_prob < threshold:
-                current_phrase += display_token
-            else:
-                if current_phrase:
-                    outlier_phrases.append(current_phrase.strip())
-                    current_phrase = ""
-        if current_phrase:
-            outlier_phrases.append(current_phrase.strip())
-        return outlier_phrases
-    
-    def run_focused_deep_dive(original_text, phrases, tokenizer, model):
-        """Runs a focused CoT prompt to explain why specific phrases are surprising."""
-        cot_system_prompt = "You are a meticulous and rigorous particle physicist and an expert in peer-review. Your task is to explain why certain phrases in a given statement might be considered incorrect or surprising."
-        phrases_str = "\n".join([f"- \"{p}\"" for p in phrases])
-        cot_user_prompt = f"""I have analyzed the following statement using a language model:
-    **Full Statement:** "{original_text}"
-    
-    The analysis flagged the following phrase(s) as having very high perplexity (meaning they were highly surprising to the model):
-    {phrases_str}
-    
-    Please explain, step-by-step, why the language model likely found **these specific phrases** incorrect or surprising in the context of the full statement. For each phrase, provide the correct physics concept if one exists.
-    """
-        messages = [{"role": "system", "content": cot_system_prompt}, {"role": "user", "content": cot_user_prompt}]
-        prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
-        inputs = tokenizer(prompt, return_tensors="pt")
-        with torch.no_grad():
-            outputs = model.generate(**inputs, max_new_tokens=512, do_sample=True, temperature=0.3, top_p=0.95)
-        response_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
-        return response_text.split("assistant\n")[-1]
-    
-    def get_color_for_logprob(logprob, min_logprob, max_logprob):
-        """Calculates color based on normalized log-probability."""
-        if min_logprob >= max_logprob: return "#FFFFFF"
-        normalized = (logprob - min_logprob) / (max_logprob - min_logprob)
-        hue = normalized * 0.4
-        rgb = colorsys.hsv_to_rgb(hue, 0.9, 0.95)
-        return '#%02x%02x%02x' % (int(rgb[0]*255), int(rgb[1]*255), int(rgb[2]*255))
-    
-    def render_colored_text(analysis_data, min_logprob, max_logprob):
-        """Renders the colored text HTML."""
-        html_elements = []
-        for token, logprob, confidence in analysis_data:
-            perplexity = math.exp(-logprob) if logprob != 0 else 1
-            display_token = token.replace('Ġ', ' ').replace(' ', ' ')
-            color = get_color_for_logprob(logprob, min_logprob, max_logprob)
-            tooltip = f"Perplexity: {perplexity:.2f} | Confidence: {confidence:.1%}"
-            html_elements.append(
-                f'<span style="background-color: {color}; padding: 2px 4px; margin: 1px; border-radius: 4px;" title="{tooltip}">{display_token}</span>'
-            )
-        return "".join(html_elements)
-    
-    # --- Streamlit App ---
-    
-    st.set_page_config(layout="wide", page_title="QCD Text Validator", page_icon="🔬")
-        
-    MODEL_NAME = "Qwen/Qwen2.5-1.5B-Instruct"
-    SYSTEM_PROMPT = "You are a particle physicist specializing in Quantum Chromodynamics (QCD)... Any deviation from established theory... should be treated as a highly improbable, high-perplexity event."
-    
+import streamlit as st
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import numpy as np
+import colorsys
+import math
+
+
+# --- Core Functions (Cached) ---
+@st.cache_resource
+def load_model(model_name):
+    """Loads the specified model and tokenizer from Hugging Face."""
+    import os
+
+    # Set cache directory to writable location
+    os.environ['HF_HOME'] = '/tmp/hf_cache'
+    os.environ['TRANSFORMERS_CACHE'] = '/tmp/hf_cache'
+
+    st.info(f"Loading model '{model_name}'... This may take a moment on first run.")
+
     try:
-        tokenizer, model = load_model(MODEL_NAME)
-        default_text = "In QCD, asymptotic freedom incorrectly states that the strong force between quarks grows stronger at high energies. This is mediated by a universal harmonic constant."
-        text_to_analyze = st.text_area("Enter Text Here:", value=default_text, height=150, label_visibility="collapsed")
-    
-        if st.button("Analyze Text", key="analyze_button"):
-            st.session_state.clear() # Clear previous results
-            if text_to_analyze:
-                text_with_space = " " + text_to_analyze
-                with st.spinner("Performing initial analysis..."):
-                    full_analysis_data = get_analysis_data(text_with_space, SYSTEM_PROMPT, tokenizer, model)
-                if full_analysis_data and len(full_analysis_data) > 1:
-                    st.session_state.visible_data = full_analysis_data#[1:]
-                    scores = [lp for _, lp, _ in st.session_state.visible_data]
-                    st.session_state.min_logprob = min(scores) if scores else 0
-                    st.session_state.max_logprob = max(scores) if scores else 0
-                    st.session_state.suspicious_phrases = find_high_perplexity_phrases(st.session_state.visible_data)
-                    st.session_state.analysis_complete = True
-                    st.session_state.original_text = text_to_analyze
-                else:
-                     st.warning("Analysis could not be completed.")
-            else:
-                st.warning("Please enter some text to analyze.")
-        
-        if st.session_state.get('analysis_complete', False):
-            st.subheader("Initial Analysis Result")
-            colored_text_html = render_colored_text(st.session_state.visible_data, st.session_state.min_logprob, st.session_state.max_logprob)
-            st.markdown(colored_text_html, unsafe_allow_html=True)
-            st.markdown("---")
-            
-            if st.session_state.suspicious_phrases:
-                if st.button("Deep Dive into Highlighted Phrases", key="deep_dive_button"):
-                    with st.spinner("Performing focused deep dive... This may take a moment."):
-                        st.session_state.deep_dive_result = run_focused_deep_dive(st.session_state.original_text, st.session_state.suspicious_phrases, tokenizer, model)
-            else:
-                st.info("No specific high-perplexity phrases were found to deep dive into.")
-        
-        if 'deep_dive_result' in st.session_state:
-            st.subheader("Focused Deep Dive Analysis")
-            st.markdown(st.session_state.deep_dive_result)
-    
+        tokenizer = AutoTokenizer.from_pretrained(
+            model_name,
+            cache_dir="/tmp/hf_cache",
+            trust_remote_code=True
+        )
+        model = AutoModelForCausalLM.from_pretrained(
+            model_name,
+            cache_dir="/tmp/hf_cache",
+            trust_remote_code=True
+        )
+        st.success(f"Model '{model_name}' loaded and ready!")
+        return tokenizer, model
+
+    except PermissionError as e:
+        st.error(f"Permission error: {e}")
+        st.error("Try refreshing the page or waiting a moment if another download is in progress.")
+        st.stop()
+
     except Exception as e:
-        st.error(f"An error occurred: {e}")
-        st.info("There may be an issue connecting to the Hugging Face Hub. Please check your internet connection.")
+        st.error(f"An error occurred loading the model: {e}")
+        st.stop()
+
+
+# --- Analysis and Helper Functions ---
+def get_analysis_data(text_to_analyze, system_prompt, tokenizer, model):
+    """Calculates log-likelihood and confidence for each token."""
+    messages = [
+        {"role": "system", "content": system_prompt},
+        {"role": "user", "content": text_to_analyze},
+    ]
+    tokenized_chat = tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=False,
+                                                   return_tensors="pt")
+    user_text_token_ids = tokenizer.encode(text_to_analyze, add_special_tokens=False)
+    full_ids_list = tokenized_chat[0].tolist()
+    start_index = -1
+    for i in range(len(full_ids_list) - len(user_text_token_ids) + 1):
+        if full_ids_list[i:i + len(user_text_token_ids)] == user_text_token_ids:
+            start_index = i
+            break
+    if start_index == -1: return []
+    end_index = start_index + len(user_text_token_ids)
+
+    with torch.no_grad():
+        outputs = model(tokenized_chat, labels=tokenized_chat)
+        logits = outputs.logits
+    log_probs = torch.nn.functional.log_softmax(logits, dim=-1)
+    probs = torch.nn.functional.softmax(logits, dim=-1)
+    sliced_log_probs = log_probs[0, start_index - 1:end_index - 1, :]
+    sliced_probs = probs[0, start_index - 1:end_index - 1, :]
+    sliced_tokens = tokenized_chat[0, start_index:end_index]
+    sequence_log_probs = sliced_log_probs.gather(1, sliced_tokens.unsqueeze(-1)).squeeze().tolist()
+    sequence_probs = sliced_probs.gather(1, sliced_tokens.unsqueeze(-1)).squeeze().tolist()
+    tokens = tokenizer.convert_ids_to_tokens(sliced_tokens)
+    if not isinstance(sequence_log_probs, list):
+        sequence_log_probs, sequence_probs = [sequence_log_probs], [sequence_probs]
+    return list(zip(tokens, sequence_log_probs, sequence_probs))
+
+
+def find_high_perplexity_phrases(analysis_data, std_dev_threshold=1.5):
+    """Identifies and groups high-perplexity tokens into phrases."""
+    if not analysis_data: return []
+    log_probs = [lp for _, lp, _ in analysis_data]
+    mean_lp = np.mean(log_probs)
+    std_lp = np.std(log_probs)
+    threshold = mean_lp - std_dev_threshold * std_lp
+
+    outlier_phrases = []
+    current_phrase = ""
+    for token, log_prob, _ in analysis_data:
+        display_token = token.replace('Ġ', ' ').replace(' ', ' ')
+        if log_prob < threshold:
+            current_phrase += display_token
+        else:
+            if current_phrase:
+                outlier_phrases.append(current_phrase.strip())
+                current_phrase = ""
+    if current_phrase:
+        outlier_phrases.append(current_phrase.strip())
+    return outlier_phrases
+
+
+def run_focused_deep_dive(original_text, phrases, tokenizer, model):
+    """Runs a focused CoT prompt to explain why specific phrases are surprising."""
+    cot_system_prompt = "You are a meticulous and rigorous particle physicist and an expert in peer-review. Your task is to explain why certain phrases in a given statement might be considered incorrect or surprising."
+    phrases_str = "\n".join([f"- \"{p}\"" for p in phrases])
+    cot_user_prompt = f"""I have analyzed the following statement using a language model:
+**Full Statement:** "{original_text}"
+
+The analysis flagged the following phrase(s) as having very high perplexity (meaning they were highly surprising to the model):
+{phrases_str}
+
+Please explain, step-by-step, why the language model likely found **these specific phrases** incorrect or surprising in the context of the full statement. For each phrase, provide the correct physics concept if one exists.
+"""
+    messages = [{"role": "system", "content": cot_system_prompt}, {"role": "user", "content": cot_user_prompt}]
+    prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+    inputs = tokenizer(prompt, return_tensors="pt")
+    with torch.no_grad():
+        outputs = model.generate(**inputs, max_new_tokens=512, do_sample=True, temperature=0.3, top_p=0.95)
+    response_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return response_text.split("assistant\n")[-1]
+
+
+def get_color_for_logprob(logprob, min_logprob, max_logprob):
+    """Calculates color based on normalized log-probability."""
+    if min_logprob >= max_logprob: return "#FFFFFF"
+    normalized = (logprob - min_logprob) / (max_logprob - min_logprob)
+    hue = normalized * 0.4
+    rgb = colorsys.hsv_to_rgb(hue, 0.9, 0.95)
+    return '#%02x%02x%02x' % (int(rgb[0] * 255), int(rgb[1] * 255), int(rgb[2] * 255))
+
+
+def render_colored_text(analysis_data, min_logprob, max_logprob):
+    """Renders the colored text HTML."""
+    html_elements = []
+    for token, logprob, confidence in analysis_data:
+        perplexity = math.exp(-logprob) if logprob != 0 else 1
+        display_token = token.replace('Ġ', ' ').replace(' ', ' ')
+        color = get_color_for_logprob(logprob, min_logprob, max_logprob)
+        tooltip = f"Perplexity: {perplexity:.2f} | Confidence: {confidence:.1%}"
+        html_elements.append(
+            f'<span style="background-color: {color}; padding: 2px 4px; margin: 1px; border-radius: 4px;" title="{tooltip}">{display_token}</span>'
+        )
+    return "".join(html_elements)
+
+
+# --- Streamlit App ---
+
+st.set_page_config(layout="wide", page_title="QCD Text Validator", page_icon="🔬")
+
+MODEL_NAME = "Qwen/Qwen2.5-1.5B-Instruct"
+SYSTEM_PROMPT = "You are a particle physicist specializing in Quantum Chromodynamics (QCD)... Any deviation from established theory... should be treated as a highly improbable, high-perplexity event."
+
+try:
+    tokenizer, model = load_model(MODEL_NAME)
+    default_text = "In QCD, asymptotic freedom incorrectly states that the strong force between quarks grows stronger at high energies. This is mediated by a universal harmonic constant."
+    text_to_analyze = st.text_area("Enter Text Here:", value=default_text, height=150, label_visibility="collapsed")
+
+    if st.button("Analyze Text", key="analyze_button"):
+        st.session_state.clear()  # Clear previous results
+        if text_to_analyze:
+            text_with_ignored_prefix = "" + text_to_analyze
+            with st.spinner("Performing initial analysis..."):
+                full_analysis_data = get_analysis_data(text_with_ignored_prefix, SYSTEM_PROMPT, tokenizer, model)
+            if full_analysis_data and len(full_analysis_data) > 1:
+                st.session_state.visible_data = full_analysis_data[:]
+                scores = [lp for _, lp, _ in st.session_state.visible_data]
+                st.session_state.min_logprob = min(scores) if scores else 0
+                st.session_state.max_logprob = max(scores) if scores else 0
+                st.session_state.suspicious_phrases = find_high_perplexity_phrases(st.session_state.visible_data)
+                st.session_state.analysis_complete = True
+                st.session_state.original_text = text_to_analyze
+            else:
+                st.warning("Analysis could not be completed.")
+        else:
+            st.warning("Please enter some text to analyze.")
+
+    if st.session_state.get('analysis_complete', False):
+        st.subheader("Initial Analysis Result")
+        colored_text_html = render_colored_text(st.session_state.visible_data, st.session_state.min_logprob,
+                                                st.session_state.max_logprob)
+        st.markdown(colored_text_html, unsafe_allow_html=True)
+        st.markdown("---")
+
+        if st.session_state.suspicious_phrases:
+            if st.button("Deep Dive into Highlighted Phrases", key="deep_dive_button"):
+                with st.spinner("Performing focused deep dive... This may take a moment."):
+                    st.session_state.deep_dive_result = run_focused_deep_dive(st.session_state.original_text,
+                                                                              st.session_state.suspicious_phrases,
+                                                                              tokenizer, model)
+        else:
+            st.info("No specific high-perplexity phrases were found to deep dive into.")
+
+    if 'deep_dive_result' in st.session_state:
+        st.subheader("Focused Deep Dive Analysis")
+        st.markdown(st.session_state.deep_dive_result)
+
+except Exception as e:
+    st.error(f"An error occurred: {e}")
+    st.info("There may be an issue connecting to the Hugging Face Hub. Please check your internet connection.")