Spaces:
Sleeping
Sleeping
File size: 1,654 Bytes
b0bb2e9 | 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 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 | from fastapi import FastAPI
from pydantic import BaseModel
from transformers import DistilBertTokenizerFast, DistilBertForSequenceClassification
import torch
import torch.nn.functional as F
# =========================
# INIT
# =========================
app = FastAPI(
title="Skill Classification API",
description="Predicts skill from student check-ins",
version="1.0"
)
MODEL_PATH = "mjpsm/skill-classifier-BERT-v1"
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print("🔄 Loading model...")
tokenizer = DistilBertTokenizerFast.from_pretrained(MODEL_PATH)
model = DistilBertForSequenceClassification.from_pretrained(MODEL_PATH)
model.to(device)
model.eval()
print("✅ Model loaded!")
# =========================
# INPUT SCHEMA
# =========================
class InputText(BaseModel):
text: str
# =========================
# ROOT
# =========================
@app.get("/")
def home():
return {"message": "Skill Classification API is running"}
# =========================
# PREDICT
# =========================
@app.post("/predict")
def predict(input: InputText):
text = input.text
inputs = tokenizer(
text,
return_tensors="pt",
truncation=True,
padding=True,
max_length=128
)
inputs = {k: v.to(device) for k, v in inputs.items()}
with torch.no_grad():
outputs = model(**inputs)
probs = F.softmax(outputs.logits, dim=1)
pred = torch.argmax(probs, dim=1).item()
label = model.config.id2label[pred]
confidence = probs[0][pred].item()
return {
"prediction": label,
"confidence": confidence
} |