Spaces:
Sleeping
Sleeping
File size: 1,649 Bytes
52c9aac |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 |
import torch
from transformers import AutoTokenizer, AutoModelForMaskedLM
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
import uvicorn
# Initialize FastAPI
app = FastAPI()
# --- MODEL LOADING (Runs once at startup) ---
MODEL_NAME = "naver/splade-cocondenser-ensembledistil"
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Loading model on: {device}...")
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForMaskedLM.from_pretrained(MODEL_NAME).to(device)
model.eval()
print("Model loaded successfully.")
# Input Schema
class TextRequest(BaseModel):
text: str
@app.get("/")
def home():
return {"status": "SPLADE API is running", "device": device}
@app.post("/splade")
@torch.no_grad()
def get_splade_vector(request: TextRequest):
text = request.text
if not text.strip():
raise HTTPException(status_code=400, detail="Input text cannot be empty.")
# Tokenize and move to GPU
inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=512).to(device)
# Inference
logits = model(**inputs).logits # [1, seq_len, vocab_size]
# SPLADE Logic: log(1 + ReLU(logits)) + max-pooling
term_scores = torch.log1p(torch.relu(logits))
term_importance = term_scores.max(dim=1).values.squeeze(0) # [vocab_size]
# Extract non-zero values
nz = torch.nonzero(term_importance, as_tuple=True)[0]
weights = term_importance[nz]
# Convert to standard Python lists (CPU)
indices = nz.cpu().tolist()
values = weights.cpu().float().tolist()
return {"indices": indices, "values": values} |