Update model eval

Nested/nn/BertSeqTagger.py

@@ -0,0 +1,14 @@
+import torch.nn as nn
+from transformers import BertModel
+
+class BertSeqTagger(nn.Module):
+    def __init__(self, bert_model, num_labels=2, dropout=0.1):
+        super().__init__()
+        self.bert = BertModel.from_pretrained(bert_model)
+        self.dropout = nn.Dropout(dropout)
+        self.linear = nn.Linear(768, num_labels)
+    def forward(self, x):
+        y = self.bert(x)
+        y = self.dropout(y["last_hidden_state"])
+        logits = self.linear(y)
+        return logits

app.py

@@ -2,10 +2,8 @@ from fastapi import FastAPI
 import torch
 import pickle
 from huggingface_hub import hf_hub_download
+from Nested.nn.BertSeqTagger import BertSeqTagger
 
-import os
-
-print(os.getcwd())
 app = FastAPI()
 print("Version 2...")
 
@@ -24,5 +22,54 @@ checkpoint_path = hf_hub_download(
 with open("Nested/utils/tag_vocab.pkl", "rb") as f:
     id2label = pickle.load(f)
 
-model = torch.load(checkpoint_path, map_location="cpu")
+# model = torch.load(checkpoint_path, map_location="cpu")
+model = BertSeqTagger(
+    pretrained_path="aubmindlab/bert-base-arabertv2",
+    dropout_p=0.1
+)
+
+
+def load_model_from_checkpoint(model, checkpoint, strict=True):
+    if isinstance(checkpoint, torch.nn.Module):
+        return checkpoint
+
+    if not isinstance(checkpoint, dict):
+        raise TypeError(f"Unsupported checkpoint type: {type(checkpoint)}")
+
+    candidates = [
+        "state_dict",
+        "model_state_dict",
+        "model",
+        "net",
+        "network",
+        "model_state",
+    ]
+
+    state_dict = None
+    for k in candidates:
+        if k in checkpoint and isinstance(checkpoint[k], dict):
+            state_dict = checkpoint[k]
+            break
+
+    if state_dict is None:
+        looks_like_state = (
+            len(checkpoint) > 0
+            and all(isinstance(v, torch.Tensor) for v in checkpoint.values())
+            and all(isinstance(k, str) for k in checkpoint.keys())
+        )
+        if looks_like_state:
+            state_dict = checkpoint
+        else:
+            raise KeyError(f"No model weights found. Keys: {list(checkpoint.keys())}")
+
+    if len(state_dict) > 0:
+        any_key = next(iter(state_dict.keys()))
+        if any_key.startswith("module."):
+            state_dict = {k.replace("module.", "", 1): v for k, v in state_dict.items()}
+
+    model.load_state_dict(state_dict, strict=strict)
+    return model
+
+ckpt = torch.load(checkpoint_path, map_location="cpu")
+model = load_model_from_checkpoint(model, ckpt, strict=False)
 model.eval()