dhindman/ltc25-vision-statement-categorizer

1from fastapi import FastAPI2from fastapi.staticfiles import StaticFiles3from fastapi.responses import FileResponse, JSONResponse4from fastapi.middleware.cors import CORSMiddleware5from pydantic import BaseModel6from sentence_transformers import SentenceTransformer, util7from typing import Dict, List, Optional8import json9import torch10 11# Initialize the FastAPI app12app = FastAPI()13 14# Mount the static directory to serve the HTML file15app.mount("/static", StaticFiles(directory="static"), name="static")16 17# Enable CORS18app.add_middleware(19    CORSMiddleware,20    allow_origins=["*"],  # Allows all origins21    allow_credentials=True,22    allow_methods=["*"],  # Allows all methods23    allow_headers=["*"],  # Allows all headers24)25 26# Load the sentence embedding model27model = SentenceTransformer('intfloat/e5-large-v2')28 29# --- Caching Mechanism ---30embedding_cache: Dict[str, torch.Tensor] = {}31 32# Load categories from the JSON file33with open("categories.json", "r", encoding="utf-8") as f:34    categories_from_file = json.load(f)35 36# Pre-compute and cache the embeddings for the example statements from the file37category_embeddings_from_file: Dict[str, torch.Tensor] = {}38for category, examples in categories_from_file.items():39    prefixed_examples = ["passage: " + example for example in examples]40    embeddings = model.encode(prefixed_examples, convert_to_tensor=True)41    category_embeddings_from_file[category] = embeddings42    # Populate the cache43    for i, example in enumerate(prefixed_examples):44        embedding_cache[example] = embeddings[i]45# --- End Caching Mechanism ---46 47class MatchRequest(BaseModel):48    statement: str49    categories: Optional[Dict[str, List[str]]] = None50 51@app.get("/")52async def read_root():53    return FileResponse('static/index.html')54 55@app.get("/categories")56async def get_categories():57    """58    Returns the categories and example statements from categories.json.59    """60    return JSONResponse(content=categories_from_file)61 62@app.post("/match")63async def match_statement(request: MatchRequest):64    """65    Matches a mission statement to the most closely aligned category.66    Accepts an optional 'categories' object for real-time experimentation,67    using a cache to speed up subsequent requests.68    """69    statement_embedding = model.encode("query: " + request.statement, convert_to_tensor=True)70 71    best_match = {"category": None, "score": 0.0}72 73    if request.categories:74        # If categories are provided, use the cache-aware logic75        for category, examples in request.categories.items():76            if not examples:77                continue78            79            embeddings_to_use = []80            for example in examples:81                prefixed_example = "passage: " + example82                if prefixed_example in embedding_cache:83                    # Cache hit84                    embeddings_to_use.append(embedding_cache[prefixed_example])85                else:86                    # Cache miss: compute and cache the new embedding87                    new_embedding = model.encode(prefixed_example, convert_to_tensor=True)88                    embedding_cache[prefixed_example] = new_embedding89                    embeddings_to_use.append(new_embedding)90            91            if not embeddings_to_use:92                continue93 94            category_embeddings = torch.stack(embeddings_to_use)95            similarities = util.pytorch_cos_sim(statement_embedding, category_embeddings)96            max_similarity = similarities.max().item()97            if max_similarity > best_match["score"]:98                best_match["category"] = category99                best_match["score"] = max_similarity100    else:101        # Otherwise, use the fully pre-computed embeddings from the file102        for category, embeddings in category_embeddings_from_file.items():103            similarities = util.pytorch_cos_sim(statement_embedding, embeddings)104            max_similarity = similarities.max().item()105            if max_similarity > best_match["score"]:106                best_match["category"] = category107                best_match["score"] = max_similarity108 109    return best_match110