Initial commit: text saliency pro with Gemma 2B

.gitignore

@@ -0,0 +1,49 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# Virtual Environments
+venv/
+env/
+.env/
+ENV/
+env.bak/
+venv.bak/
+
+# Logs and databases
+*.log
+result.png
+server.log
+
+# IDEs
+.idea/
+.vscode/
+*.swp
+*.swo
+
+# macOS
+.DS_Store

Dockerfile

@@ -0,0 +1,31 @@
+# Use official Python image
+FROM python:3.9-slim
+
+# Set the working directory
+WORKDIR /app
+
+# Install system dependencies (needed for compiling some python packages if required)
+RUN apt-get update && apt-get install -y --no-install-recommends \
+    build-essential \
+    && rm -rf /var/lib/apt/lists/*
+
+# Copy requirements
+COPY requirements.txt .
+
+# Install Python dependencies
+RUN pip install --no-cache-dir -r requirements.txt
+
+# Copy the rest of the application
+COPY . .
+
+# Hugging Face Spaces require running as a non-root user (UID 1000)
+RUN useradd -m -u 1000 user
+USER user
+ENV HOME=/home/user \
+    PATH=/home/user/.local/bin:$PATH
+
+# Expose the port HF Spaces uses
+EXPOSE 7860
+
+# Command to run the application
+CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "7860"]

main.py

@@ -0,0 +1,109 @@
+import torch
+from fastapi import FastAPI
+from fastapi.staticfiles import StaticFiles
+from fastapi.middleware.cors import CORSMiddleware
+from pydantic import BaseModel
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import uvicorn
+import os
+
+app = FastAPI()
+
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+app.mount("/static", StaticFiles(directory="static"), name="static")
+
+model_id = "google/gemma-2b"
+
+# Load the model and tokenizer globally.
+# Use MPS if available, otherwise CPU. MPS (Metal Performance Shaders) works well on modern Macs.
+device = torch.device("mps" if torch.backends.mps.is_available() else "cpu")
+print(f"Loading {model_id} on {device}...")
+
+try:
+    hf_token = os.environ.get("HF_TOKEN")
+    tokenizer = AutoTokenizer.from_pretrained(model_id, token=hf_token)
+    model = AutoModelForCausalLM.from_pretrained(
+        model_id,
+        torch_dtype=torch.bfloat16,
+        attn_implementation="eager",
+        token=hf_token
+    ).to(device)
+    print("Model loaded successfully.")
+except Exception as e:
+    print(f"Error loading model: {e}")
+    print("Make sure you are logged into Hugging Face and have access to the Gemma model.")
+    print("Run `huggingface-cli login` in your terminal.")
+
+class TextRequest(BaseModel):
+    text: str
+
+@app.post("/analyze")
+async def analyze_text(request: TextRequest):
+    text = request.text
+    if not text.strip():
+        return {"tokens": [], "scores": []}
+
+    inputs = tokenizer(text, return_tensors="pt").to(device)
+    
+    with torch.no_grad():
+        # Ensure we ask the model to output attentions explicitly
+        outputs = model(**inputs, output_attentions=True)
+        
+    # Check if attentions are actually returned
+    if not outputs.attentions:
+        print("Warning: Model did not return attentions.")
+        return {"words": []}
+        
+    # outputs.attentions is a tuple of (batch_size, num_heads, sequence_length, sequence_length)
+    # Get the last layer's attention
+    attentions = outputs.attentions[-1]
+    
+    # Average across all heads
+    avg_attention = attentions[0].mean(dim=0)  # shape: (seq_len, seq_len)
+    
+    # Calculate importance: sum of attention each token *receives* from the sequence
+    importance = avg_attention.sum(dim=0).cpu().float().numpy()
+    
+    if len(importance) > 1:
+        # Normalize to 0-1, optionally excluding the first token (<bos>) from max/min calculation
+        # as <bos> often has very high attention, skewing the rest
+        min_score = importance[1:].min()
+        max_score = importance[1:].max()
+        
+        normalized_scores = (importance - min_score) / (max_score - min_score)
+        # Keep <bos> at max score
+        normalized_scores[0] = 1.0
+        normalized_scores = normalized_scores.clip(0, 1)
+    else:
+        normalized_scores = [1.0] * len(importance)
+        
+    input_ids = inputs["input_ids"][0].tolist()
+    tokens = tokenizer.convert_ids_to_tokens(input_ids)
+    
+    result = []
+    for i, t in enumerate(tokens):
+        # Decode properly
+        word = tokenizer.decode([input_ids[i]])
+        # Special check for Gemma, decoding often removes spaces incorrectly or leaves tokens empty
+        # Let's clean the raw token just in case
+        raw_clean = t.replace('\u2581', ' ')
+        
+        # We will pass both decoded word and raw cleaned token to frontend to help render
+        result.append({
+            "token": raw_clean,
+            "word": word,
+            "score": float(normalized_scores[i])
+        })
+        
+    return {"words": result}
+
+if __name__ == "__main__":
+    port = int(os.environ.get("PORT", 7860))
+    uvicorn.run("main:app", host="0.0.0.0", port=port, reload=True)

requirements.txt

@@ -0,0 +1,6 @@
+fastapi
+uvicorn
+torch
+transformers
+pydantic
+accelerate

static/index.html

@@ -0,0 +1,228 @@
+<!DOCTYPE html>
+<html lang="en">
+<head>
+    <meta charset="UTF-8">
+    <meta name="viewport" content="width=device-width, initial-scale=1.0">
+    <title>Text Saliency Pro</title>
+    <style>
+        body {
+            font-family: system-ui, -apple-system, BlinkMacSystemFont, "Segoe UI", Roboto, "Helvetica Neue", Arial, sans-serif;
+            max-width: 800px;
+            margin: 0 auto;
+            padding: 2rem;
+            line-height: 1.5;
+            background-color: #f9fafb;
+            color: #111827;
+        }
+
+        h1 {
+            font-size: 2.5rem;
+            margin-bottom: 1rem;
+            text-align: center;
+        }
+
+        p.description {
+            text-align: center;
+            color: #4b5563;
+            margin-bottom: 2rem;
+        }
+
+        .container {
+            background: white;
+            padding: 2rem;
+            border-radius: 0.5rem;
+            box-shadow: 0 4px 6px -1px rgba(0, 0, 0, 0.1), 0 2px 4px -1px rgba(0, 0, 0, 0.06);
+        }
+
+        textarea {
+            width: 100%;
+            height: 150px;
+            padding: 0.75rem;
+            border: 1px solid #d1d5db;
+            border-radius: 0.375rem;
+            font-size: 1rem;
+            resize: vertical;
+            margin-bottom: 1rem;
+            box-sizing: border-box;
+        }
+
+        .controls {
+            display: flex;
+            align-items: center;
+            justify-content: space-between;
+            margin-bottom: 1.5rem;
+            flex-wrap: wrap;
+            gap: 1rem;
+        }
+
+        .slider-group {
+            display: flex;
+            align-items: center;
+            gap: 1rem;
+            flex-grow: 1;
+        }
+
+        input[type="range"] {
+            flex-grow: 1;
+            max-width: 300px;
+        }
+
+        button {
+            background-color: #3b82f6;
+            color: white;
+            border: none;
+            padding: 0.5rem 1.5rem;
+            font-size: 1rem;
+            border-radius: 0.375rem;
+            cursor: pointer;
+            transition: background-color 0.2s;
+        }
+
+        button:hover {
+            background-color: #2563eb;
+        }
+
+        button:disabled {
+            background-color: #9ca3af;
+            cursor: not-allowed;
+        }
+
+        #result-container {
+            margin-top: 2rem;
+            padding: 1.5rem;
+            background-color: #f3f4f6;
+            border-radius: 0.375rem;
+            min-height: 100px;
+            white-space: pre-wrap;
+            font-size: 1.125rem;
+        }
+
+        /* Token specific styles */
+        .token {
+            transition: font-weight 0.2s;
+        }
+
+        .highlighted {
+            font-weight: 800; /* Extra bold */
+            color: #000;
+        }
+        
+        #loading {
+            display: none;
+            color: #6b7280;
+            text-align: center;
+            margin-top: 1rem;
+        }
+    </style>
+</head>
+<body>
+
+    <h1>Text Saliency Pro</h1>
+    <p class="description">Improve reading comprehension using LLM attention vectors.<br>Words with attention above the threshold will be bolded.</p>
+
+    <div class="container">
+        <textarea id="text-input" placeholder="Enter or paste your text here...">In this project I want to use the attention vectors of a llm to bold the most important words in an input text to improve reading comprehension.</textarea>
+        
+        <div class="controls">
+            <button id="analyze-btn">Analyze Text</button>
+            <div class="slider-group">
+                <label for="threshold">Attention Threshold: <span id="threshold-val">0.50</span></label>
+                <input type="range" id="threshold" min="0" max="1" step="0.01" value="0.5">
+            </div>
+        </div>
+        
+        <div id="loading">Analyzing attention vectors with Gemma 2B... Please wait.</div>
+
+        <div id="result-container">
+            <!-- Processed text will appear here -->
+        </div>
+    </div>
+
+    <script>
+        const inputArea = document.getElementById('text-input');
+        const analyzeBtn = document.getElementById('analyze-btn');
+        const thresholdSlider = document.getElementById('threshold');
+        const thresholdVal = document.getElementById('threshold-val');
+        const resultContainer = document.getElementById('result-container');
+        const loading = document.getElementById('loading');
+
+        let currentTokens = []; // Array of {token: str, word: str, score: float}
+
+        // Update threshold display
+        thresholdSlider.addEventListener('input', (e) => {
+            thresholdVal.textContent = parseFloat(e.target.value).toFixed(2);
+            renderTokens(); // Re-render instantly when slider changes
+        });
+
+        // Analyze text when button is clicked
+        analyzeBtn.addEventListener('click', async () => {
+            const text = inputArea.value.trim();
+            if (!text) return;
+
+            // Update UI state
+            analyzeBtn.disabled = true;
+            loading.style.display = 'block';
+            resultContainer.innerHTML = '';
+
+            try {
+                // Call the FastAPI backend
+                const response = await fetch('/analyze', {
+                    method: 'POST',
+                    headers: {
+                        'Content-Type': 'application/json'
+                    },
+                    body: JSON.stringify({ text })
+                });
+
+                if (!response.ok) {
+                    throw new Error('Network response was not ok');
+                }
+
+                const data = await response.json();
+                currentTokens = data.words || [];
+                renderTokens();
+
+            } catch (error) {
+                console.error('Error analyzing text:', error);
+                resultContainer.innerHTML = '<span style="color: red;">Error analyzing text. Is the backend running?</span>';
+            } finally {
+                // Restore UI state
+                analyzeBtn.disabled = false;
+                loading.style.display = 'none';
+            }
+        });
+
+        // Render the tokens based on the current threshold
+        function renderTokens() {
+            if (!currentTokens.length) return;
+
+            const threshold = parseFloat(thresholdSlider.value);
+            resultContainer.innerHTML = ''; // Clear existing
+
+            currentTokens.forEach((item, index) => {
+                // Skip the <bos> token (usually first)
+                if (index === 0 && (item.token.includes('<bos>') || item.word.includes('<bos>'))) {
+                    return;
+                }
+
+                const span = document.createElement('span');
+                span.className = 'token';
+                
+                // Add bolding if score is above threshold
+                if (item.score >= threshold) {
+                    span.classList.add('highlighted');
+                }
+
+                // If the raw token started with a space, add it here.
+                // The backend replaced the special block char with a normal space.
+                // Depending on the tokenizer, 'word' might be better to display if it represents whole words,
+                // but for subwords, using the raw token with correct spacing is usually best.
+                let displayText = item.token;
+                
+                span.textContent = displayText;
+                resultContainer.appendChild(span);
+            });
+        }
+    </script>
+</body>
+</html>

test_playwright.py

@@ -0,0 +1,56 @@
+from playwright.sync_api import sync_playwright
+import time
+
+def test_app():
+    with sync_playwright() as p:
+        print("Launching browser...")
+        browser = p.chromium.launch(headless=True)
+        page = browser.new_page()
+        
+        url = "http://localhost:8000/static/index.html"
+        print(f"Navigating to {url}...")
+        page.goto(url)
+        
+        # We'll just use the default text already in the textarea:
+        # "In this project I want to use the attention vectors of a llm to bold the most important words in an input text to improve reading comprehension."
+        
+        print("Clicking the 'Analyze Text' button...")
+        page.click("#analyze-btn")
+        
+        print("Waiting for the analysis to finish (this might take a few seconds)...")
+        # Wait for the loading text to disappear and spans to appear
+        page.wait_for_selector(".token", timeout=60000)
+        
+        # Get all tokens and their classes
+        tokens = page.query_selector_all(".token")
+        
+        print("\n--- Results ---")
+        highlighted_words = []
+        full_text = []
+        
+        for token in tokens:
+            text = token.inner_text()
+            classes = token.get_attribute("class")
+            
+            # Format output
+            if "highlighted" in classes:
+                full_text.append(f"**{text}**")
+                highlighted_words.append(text)
+            else:
+                full_text.append(text)
+                
+        print("Full output with bolded words (marked by **):")
+        # Simple join (there might be spaces in the tokens themselves based on Gemma's tokenizer)
+        print("".join(full_text))
+        
+        print("\nWords that crossed the attention threshold:")
+        print(highlighted_words)
+        
+        print("\nSaving screenshot to result.png...")
+        page.screenshot(path="result.png", full_page=True)
+        
+        browser.close()
+        print("Done!")
+
+if __name__ == "__main__":
+    test_app()