sriramdev/CodeBERT_API_Space

1from fastapi import FastAPI2from pydantic import BaseModel3from transformers import RobertaTokenizer, RobertaModel4import torch5import numpy as np6 7app = FastAPI()8 9DEVICE = "cpu"10CODEBERT_MODEL = "microsoft/codebert-base"11 12 13def mean_pool(outputs, attention_mask):14    token_embeds = outputs.last_hidden_state15    mask = attention_mask.unsqueeze(-1).expand(token_embeds.size()).float()16    summed = torch.sum(token_embeds * mask, dim=1)17    counts = torch.clamp(mask.sum(dim=1), min=1e-9)18    return summed / counts19 20 21print("Loading CodeBERT model...")22 23tokenizer = RobertaTokenizer.from_pretrained(CODEBERT_MODEL)24model = RobertaModel.from_pretrained(CODEBERT_MODEL).to(DEVICE)25model.eval()26 27print("CodeBERT loaded successfully!")28 29 30class EmbedRequest(BaseModel):31    text: str32 33 34@app.get("/")35def home():36    return {37        "status": "running",38        "model": CODEBERT_MODEL39    }40 41 42@app.post("/embed")43def embed(req: EmbedRequest):44 45    inputs = tokenizer(46        req.text,47        padding=True,48        truncation=True,49        max_length=256,50        return_tensors="pt"51    ).to(DEVICE)52 53    with torch.no_grad():54        outputs = model(**inputs)55 56    embedding = mean_pool(57        outputs,58        inputs["attention_mask"]59    ).cpu().numpy()[0]60 61    return {62        "embedding": embedding.tolist()63    }