Create main.py

main.py

@@ -0,0 +1,50 @@
+from fastapi import FastAPI
+from pydantic import BaseModel
+from transformers import BertTokenizerFast, BertModel
+import torch
+import torch.nn as nn
+import os
+
+# Define constants
+MODEL_PATH = os.path.join(os.path.dirname(__file__), "model")
+WEIGHTS_PATH = os.path.join(MODEL_PATH, "bert-multilabel-model.pth")
+NUM_LABELS = 6  # Adjust based on your dataset
+
+# Initialize FastAPI app
+app = FastAPI()
+
+# Load tokenizer from local directory
+tokenizer = BertTokenizerFast.from_pretrained(MODEL_PATH)
+
+# Define the BERT-based multi-label classifier
+class BertMultiLabelClassifier(nn.Module):
+    def __init__(self):
+        super(BertMultiLabelClassifier, self).__init__()
+        self.bert = BertModel.from_pretrained(MODEL_PATH)
+        self.classifier = nn.Linear(self.bert.config.hidden_size, NUM_LABELS)
+
+    def forward(self, input_ids, attention_mask):
+        output = self.bert(input_ids=input_ids, attention_mask=attention_mask)
+        cls_output = output.last_hidden_state[:, 0, :]
+        return self.classifier(cls_output)
+
+# Load the model weights
+model = BertMultiLabelClassifier()
+model.load_state_dict(torch.load(WEIGHTS_PATH, map_location="cpu"))
+model.eval()
+
+# Input schema for prediction
+class PredictRequest(BaseModel):
+    text: str
+
+@app.get("/")
+def read_root():
+    return {"message": "Multi-label BERT model is running!"}
+
+@app.post("/predict")
+def predict(request: PredictRequest):
+    inputs = tokenizer(request.text, return_tensors="pt", truncation=True, padding=True, max_length=512)
+    with torch.no_grad():
+        logits = model(**inputs)
+        probs = torch.sigmoid(logits).squeeze().tolist()
+    return {"probabilities": probs}