src/app.py

import streamlit as st
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
import numpy as np
import colorsys
import math
import os
import streamlit.components.v1 as components
import requests 

BIN_ID = os.getenv("BIN_ID")
API_KEY = os.getenv("API_KEY")


def save_feedback(feedback_text):
    url = f"https://api.jsonbin.io/v3/b/{BIN_ID}"
    headers = {"X-Master-Key": API_KEY}
    data = {"feedback": feedback_text, "timestamp": "now"}
    
    requests.put(url, json=data, headers=headers)

MODEL_NAME = "Qwen/Qwen2.5-1.5B-Instruct"
SYSTEM_PROMPT = """System Framework: High Energy Physics Validation.

This system evaluates all statements against the experimentally-verified Standard Model (SM) of particle physics. The baseline for physical reality is defined by Quantum Field Theory as applied through the SM's SU(3) x SU(2) x U(1) gauge structure.

The following principles are foundational and non-negotiable:

1.  **Particle Content:** The fundamental constituents of matter and forces are exclusively the known quarks, leptons, gauge bosons, and the single Higgs boson of the Standard Model. Postulating other particles as established fact is a primary anomaly.

2.  **Force and Mass Generation:** The electromagnetic and weak forces are unified within the electroweak theory. Fundamental particle masses arise solely through their coupling to the Higgs field via electroweak symmetry breaking. Presenting alternative mechanisms as fact is a primary anomaly.

3.  **Conservation Laws:** All particle interactions must strictly conserve energy-momentum, charge, and lepton/baryon numbers. Statements describing interactions that violate these conservation laws are primary anomalies.

4.  **Beyond the Standard Model (BSM):** Concepts outside the SM (e.g., supersymmetry, string theory, dark matter candidates, grand unification) must be framed as theoretical, hypothetical, or speculative. Presenting any BSM concept as experimentally confirmed is a primary anomaly.

The system's function is to measure the perplexity of statements against this SM framework. Factual inaccuracies, misrepresentations, or the conflation of hypothetical concepts with established physics will register as high-perplexity tokens.
"""

def filter_token(token):
    token = token.replace('Ġ', ' ')
    return token.replace('Ċ', ' ')

# --- Core Functions (Cached) ---
@st.cache_resource
def load_model(model_name):
    """Loads the specified model and tokenizer from Hugging Face."""
    # This function now only loads the model and tokenizer without displaying status here.
    cache_dir = '/tmp/hf_cache'
    os.environ['HF_HOME'] = cache_dir
    os.environ['TRANSFORMERS_CACHE'] = cache_dir

    try:
        tokenizer = AutoTokenizer.from_pretrained(model_name, cache_dir=cache_dir, trust_remote_code=True)
        model = AutoModelForCausalLM.from_pretrained(model_name, cache_dir=cache_dir, trust_remote_code=True,
                                                     attn_implementation="eager")
        return tokenizer, model
    except Exception as e:
        # We will handle the error display in the main app body
        st.session_state.model_error = e
        return None, None


# --- Analysis and Helper Functions ---
def get_analysis_data(text_to_analyze, system_prompt, tokenizer, model):
    """Calculates log-likelihood and also returns attention data and token indices."""
    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 [], None, None, -1, -1
    end_index = start_index + len(user_text_token_ids)

    with torch.no_grad():
        outputs = model(tokenized_chat, labels=tokenized_chat, output_attentions=True)
        logits = outputs.logits
        last_layer_attention = outputs.attentions[-1][0].mean(dim=0).cpu().numpy()

    log_probs = torch.nn.functional.log_softmax(logits, dim=-1)
    sliced_log_probs = log_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()
    tokens = tokenizer.convert_ids_to_tokens(sliced_tokens)

    if not isinstance(sequence_log_probs, list):
        sequence_log_probs = [sequence_log_probs]

    full_tokens = tokenizer.convert_ids_to_tokens(tokenized_chat[0])
    return list(zip(tokens, sequence_log_probs)), full_tokens, last_layer_attention, start_index, end_index


def get_outlier_indices(analysis_data, threshold=-2.5)  :
    """
    Identifies outlier token indices using Median Absolute Deviation (MAD).
    The threshold is now more sensitive by default.
    """
    if not analysis_data or len(analysis_data) < 5:
        return np.array([])

    log_probs = np.array([lp for _, lp in analysis_data])
    median_lp = np.median(log_probs)
    mad = np.median(np.abs(log_probs - median_lp))

    if mad == 0:
        return np.array([])

    modified_z_scores = 0.6745 * (log_probs - median_lp) / mad
    return np.where(modified_z_scores < threshold)[0]

def find_high_perplexity_phrases(analysis_data, outlier_indices):
    """Groups contiguous outlier tokens into phrases."""
    if not analysis_data or outlier_indices.size == 0:
        return []

    outlier_phrases = []
    current_phrase = ""
    for i, (token, _) in enumerate(analysis_data):
        display_token = filter_token(token)
        if i in outlier_indices:
            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):
    cot_system_prompt = "You are a meticulous and rigorous particle physicist..."
    phrases_str = "\n".join([f"- \"{p}\"" for p in phrases])
    cot_user_prompt = f"""I have analyzed the following statement:
**Full Statement:** "{original_text}"
The analysis flagged these phrases as surprising:
{phrases_str}
Explain, step-by-step, why the model found **these specific phrases** surprising...
"""
    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(logprob, min_lp, max_lp, scheme='green_yellow'):
    """Generates a color based on the specified color scheme."""
    if min_lp >= max_lp:
        hue = 0.33 if scheme == 'green_yellow' else 0.0  # Default to green or red
    else:
        normalized = (logprob - min_lp) / (max_lp - min_lp)
        if scheme == 'green_yellow':
            # Scale from Yellow (0.17) to Green (0.33)
            # Higher logprob (less surprise) = greener
            hue = 0.17 + normalized * (0.33 - 0.17)
        else:  # 'yellow_red'
            # Scale from Red (0.0) to Yellow (0.17)
            # Higher logprob (less surprise) = more yellow
            hue = 0.0 + normalized * 0.17

    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, outlier_indices):
    """
    Renders text with conditional color schemes.
    - No outliers: Green-to-yellow scale for all text.
    - With outliers: Green-to-yellow for normal text, Yellow-to-red for outliers.
    """
    html_elements = []
    log_probs = np.array([lp for _, lp in analysis_data])

    if not outlier_indices.any():
        # No outliers: Use a single green-yellow scale for all tokens
        min_lp, max_lp = (log_probs.min(), log_probs.max()) if log_probs.size > 0 else (0, 0)
        for token, logprob in analysis_data:
            color = get_color(logprob, min_lp, max_lp, 'green_yellow')
            display_token = token.replace('Ġ', ' ')
            perplexity = math.exp(-logprob) if logprob != 0 else 1
            tooltip = f"Perplexity: {perplexity:.2f}"
            html_elements.append(
                f'<span style="background-color: {color}; padding: 2px 1px; margin: 0px; border-radius: 3px;" title="{tooltip}">{display_token}</span>')
    else:
        # Outliers exist: Use two different color scales
        non_outlier_mask = np.ones(len(log_probs), dtype=bool)
        non_outlier_mask[outlier_indices] = False

        non_outlier_lps = log_probs[non_outlier_mask]
        outlier_lps = log_probs[outlier_indices]

        min_non_outlier, max_non_outlier = (
        non_outlier_lps.min(), non_outlier_lps.max()) if non_outlier_lps.size > 0 else (0, 0)
        min_outlier, max_outlier = (outlier_lps.min(), outlier_lps.max()) if outlier_lps.size > 0 else (0, 0)

        for i, (token, logprob) in enumerate(analysis_data):
            display_token = token.replace('Ġ', ' ')
            perplexity = math.exp(-logprob) if logprob != 0 else 1
            tooltip = f"Perplexity: {perplexity:.2f}"
            if i in outlier_indices:
                color = get_color(logprob, min_outlier, max_outlier, 'yellow_red')
            else:
                color = get_color(logprob, min_non_outlier, max_non_outlier, 'green_yellow')
            html_elements.append(
                f'<span style="background-color: {color}; padding: 2px 1px; margin: 0px; border-radius: 3px;" title="{tooltip}">{display_token}</span>')

    return "".join(html_elements)


def render_interactive_text(tokens, attention_matrix, start_index, threshold):
    """Generates interactive HTML to highlight attention targets on hover."""
    css = """<style>.interactive-text-container{line-height:2.0;font-size:1.1em}.token{cursor:pointer;padding:2px 4px;border-radius:4px;transition:background-color .2s ease-in-out}.source-highlight{background-color:#ffd700;color:#000}.target-highlight{background-color:#1e90ff;color:#fff}</style>"""
    token_spans = []
    for i, token_text in enumerate(tokens):
        original_i = start_index + i
        display_text = token_text.replace('Ġ', ' ')
        targets = []
        for j, _ in enumerate(tokens):
            if i == j: continue
            original_j = start_index + j
            score = max(attention_matrix[original_i, original_j], attention_matrix[original_j, original_i])
            if score > threshold: targets.append(f"'token-{original_j}'")
        targets_str = ",".join(targets)
        token_id = f"token-{original_i}"
        span = f'<span class="token" id="{token_id}" onmouseover="highlightTargets(\'{token_id}\',[{targets_str}])" onmouseout="clearHighlights()">{display_text}</span>'
        token_spans.append(span)
    js = """<script>const allTokens=document.querySelectorAll('.token');function highlightTargets(e,t){clearHighlights();const n=document.getElementById(e);n&&n.classList.add('source-highlight'),t.forEach(e=>{const t=document.getElementById(e);t&&t.classList.add('target-highlight')})}function clearHighlights(){allTokens.forEach(e=>{e.classList.remove('source-highlight'),e.classList.remove('target-highlight')})}</script>"""
    html_body = f'<div class="interactive-text-container">{"".join(token_spans)}</div>'
    return f"<html><head>{css}</head><body>{html_body}{js}</body></html>"


# --- Streamlit App ---
st.set_page_config(layout="wide", page_title="QCD Text Validator & Inspector", page_icon="🔬")

# Load model and tokenizer, keeping the UI clean
if 'model' not in st.session_state:
    st.session_state.tokenizer, st.session_state.model = None, None
    st.session_state.model_status = f"Loading model '{MODEL_NAME}'..."
    tokenizer, model = load_model(MODEL_NAME)
    if model:
        st.session_state.tokenizer, st.session_state.model = tokenizer, model
        st.session_state.model_status = f"✅ Model '{MODEL_NAME}' loaded successfully."
    else:
        st.session_state.model_status = f"❌ Error loading model: {st.session_state.model_error}"

tokenizer, model = st.session_state.tokenizer, st.session_state.model

if model:
    st.markdown(f'Currently running on {MODEL_NAME}, contact [https://sulcantonin.github.io/](https://sulcantonin.github.io/), [sulc.antonin@gmail.com](mailto:sulc.antonin@gmail.com)')
    default_text = "The running of the strong coupling of QCD increases with the energy scale."
    text_to_analyze = st.text_area("Enter Text to Analyze:", value=default_text, height=150)

    if st.button("Analyze Text", key="analyze_button", type="primary"):
        for key in list(st.session_state.keys()):
            if key not in ['tokenizer', 'model', 'model_status']:
                del st.session_state[key]
        if text_to_analyze:
            save_feedback(text_to_analyze)
            with st.spinner("Performing analysis..."):
                analysis_data, full_tokens, attention_matrix, start_idx, end_idx = get_analysis_data(text_to_analyze,
                                                                                                     SYSTEM_PROMPT,
                                                                                                     tokenizer, model)
            if analysis_data:
                st.session_state.analysis_data = analysis_data
                st.session_state.full_tokens = full_tokens
                st.session_state.attention_matrix = attention_matrix
                st.session_state.start_index = start_idx
                st.session_state.end_index = end_idx
                outlier_indices = get_outlier_indices(analysis_data)
                st.session_state.outlier_indices = outlier_indices
                st.session_state.suspicious_phrases = find_high_perplexity_phrases(analysis_data, outlier_indices)
                st.session_state.original_text = text_to_analyze
                st.session_state.analysis_complete = True
            else:
                st.warning("Analysis could not be completed. The input text might be too short or unusual.")
        else:
            st.warning("Please enter some text to analyze.")

    if st.session_state.get('analysis_complete', False):
        st.markdown("---")
        st.subheader("📝 Perplexity Analysis")
        st.markdown(
            "Color indicates model surprise. Green is predictable, yellow is less so. Red highlights statistical outliers.")
        colored_text_html = render_colored_text(st.session_state.analysis_data, st.session_state.outlier_indices)
        st.markdown(colored_text_html, unsafe_allow_html=True)
        st.markdown("---")


        # Attention analysis 
        # st.subheader("💡 Interactive Attention")
        # st.markdown("Hover over any word to highlight other words it pays strong attention to.")
        # start, end = st.session_state.start_index, st.session_state.end_index
        # user_tokens, user_attention_matrix = st.session_state.full_tokens[start:end], st.session_state.attention_matrix
        # max_attention = float(np.max(user_attention_matrix)) if user_attention_matrix.size > 0 else 0.1
        # default_slider_val = min(0.1, max_attention) if max_attention > 0 else 0.1
        # attention_threshold = st.slider("Attention Threshold", 0.0, max_attention, default_slider_val, 0.01, "%.2f")
        # interactive_html = render_interactive_text(user_tokens, user_attention_matrix, start, attention_threshold)
        # components.html(interactive_html, height=200, scrolling=True)
        # st.markdown("---")

        # explanatios!
        if st.session_state.suspicious_phrases:
            st.warning("High-perplexity phrases identified:")
            for phrase in st.session_state.suspicious_phrases: st.markdown(f"- *{filter_token(phrase)}*")
            if st.button("Run Deep Dive on Suspicious Phrases", key="deep_dive_button"):
                with st.spinner("Performing focused deep dive..."):
                    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 statistically significant high-perplexity phrases were found.")

        if 'deep_dive_result' in st.session_state:
            st.subheader("🧠 Focused Deep Dive Analysis")
            st.markdown(st.session_state.deep_dive_result)

# Display model status at the bottom
with st.expander("System Status", expanded=False):
    st.info(st.session_state.get('model_status', 'Initializing...'))