Upload 8 files

Dockerfile

@@ -0,0 +1,13 @@
+FROM python:3.10-slim
+
+WORKDIR /app
+
+COPY ./app /app/app
+COPY ./tokenizer /app/tokenizer
+COPY ./saved_models /app/saved_models
+COPY label_encoders.pkl /app/
+COPY app/requirements.txt /app/
+
+RUN pip install --no-cache-dir -r /app/requirements.txt
+
+CMD ["uvicorn", "app.main:app", "--host", "0.0.0.0", "--port", "7860"]

ROBERTA_model.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:61d5f8d8c55420afc000f5033303795d1cf2544f451dac80b878e2054899b4bb
+size 18

config.py

@@ -0,0 +1,21 @@
+import torch
+import os
+
+DATA_PATH = '/app/synthetic_transactions_samples_5000.csv'
+TOKENIZER_PATH = './tokenizer/'
+LABEL_ENCODERS_PATH = './label_encoders.pkl'
+MODEL_SAVE_DIR = './saved_models/'
+PREDICTIONS_SAVE_DIR = './predictions/'
+
+TEXT_COLUMN = "Sanction_Context"
+LABEL_COLUMNS = [
+    "Red_Flag_Reason",
+    "Maker_Action",
+    "Escalation_Level",
+    "Risk_Category",
+    "Risk_Drivers",
+    "Investigation_Outcome"
+]
+MAX_LEN = 128
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+ROBERTA_MODEL_NAME = 'roberta-base'

dataset_utils.py

@@ -0,0 +1,5 @@
+import pickle
+
+def load_label_encoders(path='./label_encoders.pkl'):
+    with open(path, 'rb') as f:
+        return pickle.load(f)

docker-compose.yml

@@ -0,0 +1,14 @@
+version: '3.9'
+
+services:
+  roberta-api:
+    build:
+      context: .
+      dockerfile: Dockerfile
+    ports:
+      - "7860:7860"
+    volumes:
+      - ./app:/app/app
+      - ./saved_models:/app/saved_models
+      - ./tokenizer:/app/tokenizer
+      - ./label_encoders.pkl:/app/label_encoders.pkl

main.py

@@ -0,0 +1,40 @@
+from fastapi import FastAPI
+from pydantic import BaseModel
+import torch
+from transformers import RobertaTokenizer
+from app.roberta_model import RobertaMultiOutputModel
+from app.dataset_utils import load_label_encoders
+from app.config import MAX_LEN, LABEL_COLUMNS, MODEL_SAVE_DIR, LABEL_ENCODERS_PATH, TOKENIZER_PATH
+
+app = FastAPI()
+
+class InputText(BaseModel):
+    sanction_context: str
+
+label_encoders = load_label_encoders(LABEL_ENCODERS_PATH)
+num_classes_per_label = [len(label_encoders[col].classes_) for col in LABEL_COLUMNS]
+
+model = RobertaMultiOutputModel(num_classes_per_label)
+model.load_state_dict(torch.load(f"{MODEL_SAVE_DIR}/ROBERTA_model.pth", map_location="cpu"))
+model.eval()
+
+tokenizer = RobertaTokenizer.from_pretrained(TOKENIZER_PATH)
+
+@app.post("/predict")
+def predict(input_text: InputText):
+    inputs = tokenizer(
+        input_text.sanction_context,
+        padding='max_length',
+        truncation=True,
+        max_length=MAX_LEN,
+        return_tensors="pt"
+    )
+    with torch.no_grad():
+        outputs = model(**inputs)
+        predicted = [torch.argmax(logit, dim=1).item() for logit in outputs]
+
+    decoded = {
+        label: label_encoders[label].inverse_transform([pred])[0]
+        for label, pred in zip(LABEL_COLUMNS, predicted)
+    }
+    return {"predictions": decoded}

requirements.txt

@@ -0,0 +1,9 @@
+fastapi
+uvicorn
+transformers>=4.41.2
+torch>=2.0.0
+scikit-learn
+pandas
+numpy
+tqdm
+regex

roberta_model.py

@@ -0,0 +1,19 @@
+import torch
+import torch.nn as nn
+from transformers import RobertaModel
+from app.config import ROBERTA_MODEL_NAME
+
+class RobertaMultiOutputModel(nn.Module):
+    def __init__(self, num_labels):
+        super(RobertaMultiOutputModel, self).__init__()
+        self.roberta = RobertaModel.from_pretrained(ROBERTA_MODEL_NAME)
+        self.dropout = nn.Dropout(0.3)
+        self.classifiers = nn.ModuleList([
+            nn.Linear(self.roberta.config.hidden_size, n_classes) for n_classes in num_labels
+        ])
+
+    def forward(self, input_ids, attention_mask):
+        output = self.roberta(input_ids=input_ids, attention_mask=attention_mask)
+        pooled_output = output.pooler_output
+        pooled_output = self.dropout(pooled_output)
+        return [classifier(pooled_output) for classifier in self.classifiers]