Add inference handler for model deployment

handler.py

@@ -0,0 +1,107 @@
+
+import torch
+import torch.nn as nn
+import numpy as np
+from transformers import AutoTokenizer, AutoModel
+from torchcrf import CRF
+
+class PositionalEncoding(nn.Module):
+    def __init__(self, d_model, dropout=0.1, max_len=5000):
+        super(PositionalEncoding, self).__init__()
+        self.dropout = nn.Dropout(p=dropout)
+        pe = torch.zeros(max_len, d_model)
+        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
+        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-np.log(10000.0) / d_model))
+        pe[:, 0::2] = torch.sin(position * div_term)
+        pe[:, 1::2] = torch.cos(position * div_term)
+        pe = pe.unsqueeze(0)
+        self.register_buffer('pe', pe)
+    
+    def forward(self, x):
+        x = x + self.pe[:, :x.size(1), :]
+        return self.dropout(x)
+
+class VanillaTransformer(nn.Module):
+    def __init__(self, d_model=768, nhead=8, num_layers=3, dim_feedforward=2048, dropout=0.1):
+        super(VanillaTransformer, self).__init__()
+        self.pos_encoder = PositionalEncoding(d_model, dropout)
+        encoder_layer = nn.TransformerEncoderLayer(
+            d_model=d_model, nhead=nhead, dim_feedforward=dim_feedforward,
+            dropout=dropout, activation='gelu', batch_first=True
+        )
+        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
+        self.d_model = d_model
+    
+    def forward(self, src, src_mask=None, src_key_padding_mask=None):
+        src = self.pos_encoder(src)
+        output = self.transformer(src, mask=src_mask, src_key_padding_mask=src_key_padding_mask)
+        return output
+
+class HierarchicalLegalSegModel(nn.Module):
+    def __init__(self, longformer_model, num_labels, hidden_dim=768, transformer_layers=3, transformer_heads=8, dropout=0.1):
+        super(HierarchicalLegalSegModel, self).__init__()
+        self.longformer = longformer_model
+        self.hidden_dim = hidden_dim
+        self.vanilla_transformer = VanillaTransformer(d_model=hidden_dim, nhead=transformer_heads, num_layers=transformer_layers, dim_feedforward=hidden_dim*4, dropout=dropout)
+        self.classifier = nn.Linear(hidden_dim, num_labels)
+        self.crf = CRF(num_labels, batch_first=True)
+        self.dropout = nn.Dropout(dropout)
+        self.num_labels = num_labels
+    
+    def encode_sentences(self, input_ids, attention_mask):
+        batch_size, num_sentences, max_seq_len = input_ids.shape
+        input_ids_flat = input_ids.view(-1, max_seq_len)
+        attention_mask_flat = attention_mask.view(-1, max_seq_len)
+        outputs = self.longformer(input_ids=input_ids_flat, attention_mask=attention_mask_flat)
+        cls_embeddings = outputs.last_hidden_state[:, 0, :]
+        sentence_embeddings = cls_embeddings.view(batch_size, num_sentences, self.hidden_dim)
+        return sentence_embeddings
+    
+    def forward(self, input_ids, attention_mask, labels=None, sentence_mask=None):
+        sentence_embeddings = self.encode_sentences(input_ids, attention_mask)
+        sentence_embeddings = self.dropout(sentence_embeddings)
+        transformer_output = self.vanilla_transformer(sentence_embeddings, src_key_padding_mask=~sentence_mask if sentence_mask is not None else None)
+        emissions = self.classifier(transformer_output)
+        if labels is not None:
+            loss = -self.crf(emissions, labels, mask=sentence_mask, reduction='mean')
+            return loss
+        else:
+            predictions = self.crf.decode(emissions, mask=sentence_mask)
+            return predictions
+
+class EndpointHandler:
+    def __init__(self, path=""):
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.tokenizer = AutoTokenizer.from_pretrained(path)
+        
+        longformer = AutoModel.from_pretrained("lexlms/legal-longformer-base")
+        longformer = longformer.to(self.device)
+        for param in longformer.parameters():
+            param.requires_grad = False
+        
+        self.model = HierarchicalLegalSegModel(longformer, num_labels=7)
+        checkpoint = torch.load(f"{path}/model.pth", map_location=self.device)
+        if isinstance(checkpoint, dict) and 'model_state_dict' in checkpoint:
+            self.model.load_state_dict(checkpoint['model_state_dict'])
+        else:
+            self.model.load_state_dict(checkpoint)
+        self.model = self.model.to(self.device)
+        self.model.eval()
+        
+        self.id2label = {
+            0: "Arguments of Petitioner", 1: "Arguments of Respondent", 2: "Decision",
+            3: "Facts", 4: "Issue", 5: "None", 6: "Reasoning"
+        }
+    
+    def __call__(self, data):
+        text = data.get("inputs", "")
+        encoded = self.tokenizer(text, padding="max_length", truncation=True, max_length=512, return_tensors="pt")
+        input_ids = encoded["input_ids"].unsqueeze(1).to(self.device)
+        attention_mask = encoded["attention_mask"].unsqueeze(1).to(self.device)
+        sentence_mask = torch.ones(1, 1, dtype=torch.bool).to(self.device)
+        
+        with torch.no_grad():
+            predictions = self.model(input_ids, attention_mask, sentence_mask=sentence_mask)
+        
+        label_id = predictions[0][0]
+        return [{"label": self.id2label[label_id], "score": 1.0}]