server.py

import os
import gc
import time
import math
from typing import List, Dict, Optional
import torch
import torch.nn as nn
from transformers import AutoTokenizer, AutoModelForCausalLM
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel

app = FastAPI(title="Attention Visualizer & Token Playground")

# Enable CORS for frontend flexibility
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Global variables to cache model/tokenizer
MODEL_ID = "Qwen/Qwen2.5-0.5B-Instruct"
tokenizer = None
model = None
device = "cuda" if torch.cuda.is_available() else "cpu"

import threading
model_lock = threading.Lock()

# Custom Multihead Attention Module (Week 2 Exercise 4 logic)
class CustomMultiheadAttention(nn.Module):
    def __init__(self, embed_dim, num_heads):
        super().__init__()
        self.embed_dim = embed_dim
        self.num_heads = num_heads
        self.head_dim = embed_dim // num_heads
        self.query = nn.Linear(embed_dim, embed_dim)
        self.key = nn.Linear(embed_dim, embed_dim)
        self.value = nn.Linear(embed_dim, embed_dim)
        self.out = nn.Linear(embed_dim, embed_dim)

    def forward(self, x, mask=None):
        batch_size, seq_len, embed_dim = x.size()
        
        # Projection and head splitting
        Q = self.query(x).view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
        K = self.key(x).view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
        V = self.value(x).view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
        
        # Scaled attention score computation
        scores = torch.matmul(Q, K.transpose(-2, -1)) / (self.head_dim ** 0.5)

        if mask is not None:
            scores = scores.masked_fill(mask == 0, float("-inf"))

        # Softmax normalized attention weight matrices
        attention_weights = torch.softmax(scores, dim=-1)

        # Context-aware output projection
        attention_output = torch.matmul(attention_weights, V)
        attention_output = attention_output.transpose(1, 2).contiguous().view(batch_size, seq_len, embed_dim)
        output = self.out(attention_output)
        
        return output, attention_weights

# Instantiate global custom MHA
custom_mha = None

def lazy_load_models():
    global tokenizer, model, custom_mha
    if tokenizer is None or model is None:
        with model_lock:
            if tokenizer is None or model is None:
                print(f"Lazy loading tokenizer and model: {MODEL_ID} on device {device}...")
                tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
                if tokenizer.pad_token is None:
                    tokenizer.pad_token = tokenizer.eos_token
                    
                model_kwargs = {
                    "trust_remote_code": True,
                    "attn_implementation": "eager",
                    "torch_dtype": torch.float32 if device == "cpu" else torch.float16,
                }
                if device == "cuda":
                    model_kwargs["device_map"] = "auto"
                    
                model = AutoModelForCausalLM.from_pretrained(MODEL_ID, **model_kwargs)
                if device == "cpu":
                    model = model.to(device)
                    
                # Qwen-0.5B embed_dim is 896, num_heads is 14
                embed_dim = model.config.hidden_size
                num_heads = model.config.num_attention_heads
                print(f"Configuring Custom MHA with dim={embed_dim}, heads={num_heads}...")
                custom_mha = CustomMultiheadAttention(embed_dim=embed_dim, num_heads=num_heads).to(device)
                print("Model and Custom MHA initialized successfully!")

class AnalyzeRequest(BaseModel):
    text: str
    model_type: str = "qwen"  # "qwen" or "custom"
    layer_idx: int = 0        # 0 to 23 for Qwen

class BranchRequest(BaseModel):
    text: str
    temperature: float = 1.0
    top_k: int = 50
    top_p: float = 0.9

@app.get("/health")
def health():
    return {"status": "ok", "device": device}

@app.post("/analyze")
def analyze(req: AnalyzeRequest):
    try:
        lazy_load_models()
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Failed to load model: {str(e)}")

    if not req.text.strip():
        raise HTTPException(status_code=400, detail="Empty text input")

    inputs = tokenizer(req.text, return_tensors="pt").to(device)
    input_ids = inputs["input_ids"]
    seq_len = input_ids.shape[1]

    # Convert token IDs to individual string representations
    tokens = [tokenizer.decode([tid]) for tid in input_ids[0]]

    # Ensure layer_idx is inside valid bounds
    layer_idx = max(0, min(req.layer_idx, model.config.num_hidden_layers - 1))

    try:
        if req.model_type == "custom":
            # Pass token embeddings through custom Multihead Attention module
            with torch.no_grad():
                embeddings = model.get_input_embeddings()(input_ids)
                _, custom_attns = custom_mha(embeddings)
                # custom_attns shape: (1, num_heads, seq_len, seq_len)
                attentions_matrix = custom_attns[0].cpu().numpy().tolist()
            return {
                "tokens": tokens,
                "model_type": "custom",
                "attentions": attentions_matrix,
                "layer_idx": 0,
                "num_heads": len(attentions_matrix)
            }
        else:
            # Pass text through the real Qwen model
            with torch.no_grad():
                outputs = model(**inputs, output_attentions=True)
                # outputs.attentions is a tuple of length num_layers
                # each element is a tensor of shape: (batch, num_heads, seq_len, seq_len)
                selected_attn = outputs.attentions[layer_idx]
                attentions_matrix = selected_attn[0].float().cpu().numpy().tolist()
            return {
                "tokens": tokens,
                "model_type": "qwen",
                "attentions": attentions_matrix,
                "layer_idx": layer_idx,
                "num_heads": len(attentions_matrix)
            }
    except Exception as e:
        import traceback
        traceback.print_exc()
        raise HTTPException(status_code=500, detail=f"Inference error: {str(e)}")

@app.post("/branch")
def branch(req: BranchRequest):
    try:
        lazy_load_models()
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Failed to load model: {str(e)}")

    if not req.text.strip():
        raise HTTPException(status_code=400, detail="Empty text input")

    inputs = tokenizer(req.text, return_tensors="pt").to(device)
    
    with torch.no_grad():
        outputs = model(**inputs)
        next_token_logits = outputs.logits[0, -1, :].clone()

    # 1. Temperature scaling
    temp = max(0.01, req.temperature)
    next_token_logits = next_token_logits / temp

    # 2. Top-K filtering
    if req.top_k > 0:
        top_k = min(req.top_k, next_token_logits.shape[-1])
        indices_to_remove = next_token_logits < torch.topk(next_token_logits, top_k)[0][..., -1, None]
        next_token_logits[indices_to_remove] = float("-inf")

    # 3. Top-P (nucleus) filtering
    if req.top_p < 1.0:
        sorted_logits, sorted_indices = torch.sort(next_token_logits, descending=True)
        cumulative_probs = torch.cumsum(torch.softmax(sorted_logits, dim=-1), dim=-1)
        sorted_indices_to_remove = cumulative_probs > req.top_p
        
        # Shift to keep first token exceeding top_p
        sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
        sorted_indices_to_remove[..., 0] = 0
        indices_to_remove = sorted_indices[sorted_indices_to_remove]
        next_token_logits[indices_to_remove] = float("-inf")

    # Re-normalize probabilities
    probs = torch.softmax(next_token_logits, dim=-1)

    # Fetch top 5 options
    top_probs, top_indices = torch.topk(probs, 5)
    candidates = []
    for idx, p in zip(top_indices, top_probs):
        token_str = tokenizer.decode([idx.item()])
        candidates.append({
            "token": token_str,
            "prob": float(p.item())
        })

    # Decode recent token IDs for tracking
    prompt_tokens = [tokenizer.decode([tid]) for tid in inputs["input_ids"][0]]

    return {
        "prompt_tokens": prompt_tokens[-10:], # return last 10 tokens for visual trace
        "candidates": candidates
    }

# Check if static directory exists, mount it
static_dir = os.path.join(os.path.dirname(__file__), "static")
if not os.path.exists(static_dir):
    os.makedirs(static_dir)

# Mount files directly from visualizer folder if static folder is inside visualizer
app.mount("/", StaticFiles(directory=os.path.dirname(__file__), html=True), name="static")

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="127.0.0.1", port=8000)