Spaces:
Runtime error
Runtime error
File size: 3,124 Bytes
31954eb |
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 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 |
import os
import requests
MODEL_PATH = "/tmp/bert_model.pth"
FILE_ID = "1qqmBxbxM0CmxPGC4sqO6vLJAe-Kikiv4"
def download_from_google_drive(file_id, dest_path):
URL = "https://docs.google.com/uc?export=download"
session = requests.Session()
response = session.get(URL, params={'id': file_id}, stream=True)
def get_confirm_token(response):
for key, value in response.cookies.items():
if key.startswith('download_warning'):
return value
return None
token = get_confirm_token(response)
if token:
params = {'id': file_id, 'confirm': token}
response = session.get(URL, params=params, stream=True)
with open(dest_path, "wb") as f:
for chunk in response.iter_content(32768):
if chunk:
f.write(chunk)
if not os.path.exists(MODEL_PATH):
print("Downloading model from Google Drive...")
download_from_google_drive(FILE_ID, MODEL_PATH)
import torch
import torch.nn as nn
from transformers import BertTokenizer, BertModel
from fastapi import FastAPI
from pydantic import BaseModel
LABEL_COLUMNS = [
'Red_Flag_Reason', 'Maker_Action', 'Escalation_Level',
'Risk_Category', 'Risk_Drivers', 'Investigation_Outcome'
]
PRETRAINED_MODEL_NAME = 'bert-base-uncased'
MAX_LEN = 128
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
class BertMultiOutput(nn.Module):
def __init__(self, num_labels_per_output):
super().__init__()
self.bert = BertModel.from_pretrained(PRETRAINED_MODEL_NAME)
self.dropout = nn.Dropout(0.3)
self.classifiers = nn.ModuleList([
nn.Linear(self.bert.config.hidden_size, n_labels)
for n_labels in num_labels_per_output
])
def forward(self, input_ids, attention_mask):
outputs = self.bert(input_ids=input_ids, attention_mask=attention_mask)
pooled_output = self.dropout(outputs.pooler_output)
logits = [classifier(pooled_output) for classifier in self.classifiers]
return logits
checkpoint = torch.load(MODEL_PATH, map_location=DEVICE, weights_only=False)
label_encoders = checkpoint['label_encoders']
num_labels_list = [len(le.classes_) for le in label_encoders.values()]
model = BertMultiOutput(num_labels_list).to(DEVICE)
model.load_state_dict(checkpoint['model_state_dict'])
model.eval()
tokenizer = BertTokenizer.from_pretrained("bert_tokenizer/")
app = FastAPI()
class PredictRequest(BaseModel):
text: str
@app.get("/")
def root():
return {"message": "Multi-output BERT is ready!"}
@app.post("/predict")
def predict(request: PredictRequest):
inputs = tokenizer(
request.text,
truncation=True,
padding='max_length',
max_length=MAX_LEN,
return_tensors="pt"
).to(DEVICE)
with torch.no_grad():
outputs = model(**inputs)
preds = [torch.argmax(output, dim=1).item() for output in outputs]
decoded = {
label: label_encoders[label].inverse_transform([pred])[0]
for label, pred in zip(LABEL_COLUMNS, preds)
}
return {"predictions": decoded}
|