test

README.md

@@ -1,13 +1,40 @@
 ---
-title: Promptvisualization
-emoji: 🌖
+title: Token Probability Visualization API
+emoji: 📊
 colorFrom: blue
-colorTo: yellow
+colorTo: purple
 sdk: gradio
-sdk_version: 5.44.1
+sdk_version: 4.44.0
 app_file: app.py
 pinned: false
-short_description: api to visualize prompt engineering
+short_description: API for visualizing LLM token probabilities and embeddings
 ---
 
+# Token Probability Visualization API
+
+This Space provides API endpoints for visualizing token probabilities and embeddings from language models.
+
+## Features
+
+- **Token Probabilities**: Get next-token probability distributions for any input text
+- **Embeddings**: Access token embeddings with PCA projections for visualization
+- **Comparative Analysis**: Compare conditional vs unconditional distributions
+- **Multiple Contexts**: Analyze probability shifts across different text contexts
+
+## API Endpoints
+
+### `/predict` - Main Prediction Endpoint
+Returns comprehensive token probabilities, embeddings, and PCA projections.
+
+### `/embeddings` - Token Embeddings
+Get embeddings for specific token IDs.
+
+## Usage
+
+Use the API to build interactive visualizations showing:
+- Token probability landscapes
+- Embedding space visualizations
+- Probability distribution comparisons
+- Token movement analysis
+
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,270 @@
+import gradio as gr
+import torch
+import numpy as np
+from transformers import AutoTokenizer, AutoModelForCausalLM
+from sklearn.decomposition import PCA
+import json
+
+# Model configuration
+MODEL_ID = "openai-community/gpt2"
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+# Load model and tokenizer
+tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
+model = AutoModelForCausalLM.from_pretrained(MODEL_ID, dtype=torch.float32).to(device).eval()
+
+# Get vocabulary info
+vocab_size = tokenizer.vocab_size if tokenizer.vocab_size is not None else len(tokenizer)
+vocab_tokens = tokenizer.convert_ids_to_tokens(list(range(vocab_size)))
+
+# Cache for embeddings and PCA (computed once at startup)
+embeddings_cache = None
+pca_cache = None
+pca_projections_cache = None
+
+def initialize_embeddings():
+    """Initialize embeddings and PCA projections once at startup"""
+    global embeddings_cache, pca_cache, pca_projections_cache
+    
+    # Get embeddings
+    embeddings_cache = model.get_input_embeddings().weight.detach().cpu().numpy()  # [V, d]
+    
+    # Compute PCA
+    pca_cache = PCA(n_components=2, random_state=0)
+    pca_projections_cache = pca_cache.fit_transform(embeddings_cache)  # [V, 2]
+    
+    return embeddings_cache, pca_projections_cache
+
+# Initialize embeddings at startup
+initialize_embeddings()
+
+def nice_tok(tok: str) -> str:
+    """Clean up token display"""
+    return tok.replace("Ġ", " ").replace("▁", " ")
+
+@torch.no_grad()
+def get_next_token_probs(text: str | None):
+    """Get next token probabilities for given text context"""
+    if text and len(text) > 0:
+        enc = tokenizer(text, return_tensors="pt").to(device)
+        out = model(**enc, return_dict=True)
+        logits = out.logits[0, -1, :]
+    else:
+        # Unconditional: use BOS or EOS token
+        bos_id = tokenizer.bos_token_id or tokenizer.eos_token_id
+        if bos_id is None:
+            raise ValueError("Tokenizer has neither BOS nor EOS.")
+        input_ids = torch.tensor([[bos_id]], device=device)
+        out = model(input_ids=input_ids, return_dict=True)
+        logits = out.logits[0, -1, :]
+    
+    return logits
+
+@torch.no_grad()
+def predict_comprehensive(
+    text: str,
+    text2: str = "",
+    top_k: int = 0,
+    include_embeddings: bool = False,
+    include_pca: bool = False,
+    include_unconditional: bool = False,
+    use_logprobs: bool = True
+):
+    """
+    Comprehensive prediction endpoint that returns token probabilities,
+    embeddings, and PCA projections for visualization
+    """
+    
+    result = {
+        "model": MODEL_ID,
+        "vocab_size": vocab_size,
+        "device": device
+    }
+    
+    # Get primary text probabilities
+    logits = get_next_token_probs(text)
+    probs = torch.softmax(logits, dim=-1).cpu().numpy()
+    
+    if use_logprobs:
+        log_probs = torch.log_softmax(logits, dim=-1).cpu().numpy()
+    else:
+        log_probs = None
+    
+    # Always include vocabulary tokens (cleaned)
+    result["vocab"] = {
+        "tokens": [nice_tok(t) for t in vocab_tokens],
+        "raw_tokens": vocab_tokens,
+        "size": vocab_size
+    }
+    
+    # Primary context probabilities
+    result["probs"] = probs.tolist()
+    if log_probs is not None:
+        result["logprobs"] = log_probs.tolist()
+    
+    # Top-k tokens if requested
+    if top_k and top_k > 0:
+        vals, idxs = torch.topk(torch.from_numpy(probs), k=min(top_k, len(probs)))
+        idxs = idxs.tolist()
+        vals = vals.tolist()
+        result["topk"] = {
+            "ids": idxs,
+            "tokens": [nice_tok(vocab_tokens[i]) for i in idxs],
+            "probs": vals
+        }
+        if log_probs is not None:
+            result["topk"]["logprobs"] = [float(log_probs[i]) for i in idxs]
+    
+    # Second text context if provided
+    if text2 and len(text2) > 0:
+        logits2 = get_next_token_probs(text2)
+        probs2 = torch.softmax(logits2, dim=-1).cpu().numpy()
+        result["probs2"] = probs2.tolist()
+        
+        if use_logprobs:
+            log_probs2 = torch.log_softmax(logits2, dim=-1).cpu().numpy()
+            result["logprobs2"] = log_probs2.tolist()
+    
+    # Unconditional probabilities
+    if include_unconditional:
+        logits_uncond = get_next_token_probs(None)
+        probs_uncond = torch.softmax(logits_uncond, dim=-1).cpu().numpy()
+        result["unconditional_probs"] = probs_uncond.tolist()
+        
+        if use_logprobs:
+            log_probs_uncond = torch.log_softmax(logits_uncond, dim=-1).cpu().numpy()
+            result["unconditional_logprobs"] = log_probs_uncond.tolist()
+    
+    # Embeddings if requested
+    if include_embeddings:
+        # Return first 10 dimensions as sample (full embeddings would be too large)
+        result["embeddings_sample"] = {
+            "shape": list(embeddings_cache.shape),
+            "first_10_tokens_sample": embeddings_cache[:10, :10].tolist() if embeddings_cache is not None else None
+        }
+    
+    # PCA projections if requested
+    if include_pca:
+        if pca_projections_cache is not None:
+            result["pca"] = {
+                "projections": pca_projections_cache.tolist(),
+                "explained_variance_ratio": pca_cache.explained_variance_ratio_.tolist()
+            }
+    
+    return result
+
+@torch.no_grad()
+def get_embeddings_endpoint(token_ids: str = ""):
+    """
+    Get embeddings for specific token IDs
+    """
+    if token_ids:
+        try:
+            ids = [int(x.strip()) for x in token_ids.split(",")]
+            ids = [i for i in ids if 0 <= i < vocab_size]
+        except:
+            ids = list(range(min(100, vocab_size)))  # Default to first 100
+    else:
+        ids = list(range(min(100, vocab_size)))  # Default to first 100
+    
+    embeddings_subset = embeddings_cache[ids]
+    tokens_subset = [nice_tok(vocab_tokens[i]) for i in ids]
+    
+    return {
+        "token_ids": ids,
+        "tokens": tokens_subset,
+        "embeddings": embeddings_subset.tolist(),
+        "embedding_dim": embeddings_cache.shape[1],
+        "total_vocab_size": vocab_size
+    }
+
+# Create Gradio interface
+with gr.Blocks(title="Token Probability Visualization API") as demo:
+    gr.Markdown("""
+    # Token Probability Visualization API
+    
+    This API provides token probabilities, embeddings, and PCA projections for visualization.
+    Model: `openai-community/gpt2`
+    """)
+    
+    with gr.Tab("Comprehensive API"):
+        gr.Markdown("### Get token probabilities with optional embeddings and PCA")
+        
+        with gr.Row():
+            with gr.Column():
+                text_input = gr.Textbox(
+                    label="Primary Context",
+                    value="You are an expert in medieval history.",
+                    lines=3
+                )
+                text2_input = gr.Textbox(
+                    label="Secondary Context (optional)",
+                    placeholder="Enter second text for comparison",
+                    lines=3
+                )
+                
+                with gr.Row():
+                    top_k_slider = gr.Slider(0, 200, step=1, value=20, label="Top-K tokens")
+                    include_embeddings = gr.Checkbox(False, label="Include Embeddings Sample")
+                    include_pca = gr.Checkbox(True, label="Include PCA Projections")
+                    include_unconditional = gr.Checkbox(True, label="Include Unconditional")
+                    use_logprobs = gr.Checkbox(True, label="Include Log Probabilities")
+                
+                predict_btn = gr.Button("Get Predictions", variant="primary")
+            
+            with gr.Column():
+                output_json = gr.JSON(label="API Response")
+        
+        predict_btn.click(
+            fn=predict_comprehensive,
+            inputs=[text_input, text2_input, top_k_slider, include_embeddings, 
+                   include_pca, include_unconditional, use_logprobs],
+            outputs=output_json,
+            api_name="predict"
+        )
+    
+    with gr.Tab("Embeddings API"):
+        gr.Markdown("### Get embeddings for specific tokens")
+        
+        with gr.Row():
+            with gr.Column():
+                token_ids_input = gr.Textbox(
+                    label="Token IDs (comma-separated)",
+                    placeholder="e.g., 0,1,2,3,4 or leave empty for first 100",
+                    value="0,1,2,3,4,5,6,7,8,9"
+                )
+                get_embeddings_btn = gr.Button("Get Embeddings", variant="primary")
+            
+            with gr.Column():
+                embeddings_output = gr.JSON(label="Embeddings Response")
+        
+        get_embeddings_btn.click(
+            fn=get_embeddings_endpoint,
+            inputs=token_ids_input,
+            outputs=embeddings_output,
+            api_name="embeddings"
+        )
+    
+    gr.Markdown("""
+    ## API Usage
+    
+    ### Endpoints:
+    - `/predict`: Main endpoint for token probabilities with optional embeddings and PCA
+    - `/embeddings`: Get embeddings for specific token IDs
+    
+    ### Response includes:
+    - Token probabilities (conditional and unconditional)
+    - Log probabilities
+    - PCA projections for all vocabulary tokens
+    - Token embeddings (sample or specific)
+    - Vocabulary mappings
+    
+    ### Use this data to:
+    - Visualize token probability landscapes
+    - Compare conditional vs unconditional distributions
+    - Create embedding visualizations
+    - Build interactive token explorers
+    """)
+
+if __name__ == "__main__":
+    demo.launch(server_name="0.0.0.0", server_port=7860, share=False)

requirements.txt

@@ -0,0 +1,6 @@
+gradio==4.44.0
+torch>=2.0.0
+transformers>=4.30.0
+scikit-learn>=1.3.0
+numpy>=1.24.0
+spaces