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)