Update app file

app.py

@@ -5,6 +5,10 @@ from huggingface_hub import hf_hub_download
 from Nested.nn.BertSeqTagger import BertSeqTagger
 from transformers import AutoTokenizer, AutoModel
 import inspect
+from collections import namedtuple
+from Nested.utils.helpers import load_checkpoint
+from Nested.utils.data import get_dataloaders, text2segments
+
 app = FastAPI()
 print("Version 2...")
 
@@ -24,150 +28,45 @@ checkpoint_path = hf_hub_download(
     filename="checkpoints/checkpoint_2.pt"
 )
 
+args_path = hf_hub_download(
+    repo_id="SinaLab/Nested",
+    filename="args.json"
+)
+
 # Load model
 with open("Nested/utils/tag_vocab.pkl", "rb") as f:
     label_vocab = pickle.load(f)
 
-# model = torch.load(checkpoint_path, map_location="cpu")
-model = BertSeqTagger(
-    bert_model="aubmindlab/bert-base-arabertv2",
-    dropout=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()
-
-def predict_ner(sentence, tagger, encoder, tokenizer, id2label, device="cpu", max_length=128):
-    tagger.to(device).eval()
-    encoder.to(device).eval()
-
-    words = sentence.split()
-
-    enc = tokenizer(
-        words,
-        is_split_into_words=True,
-        return_tensors="pt",
-        truncation=True,
-        max_length=max_length
-    )
-    enc = {k: v.to(device) for k, v in enc.items()}
-
-    with torch.no_grad():
-        x = encoder(
-            input_ids=enc["input_ids"],
-            attention_mask=enc.get("attention_mask", None)
-        ).last_hidden_state  # [1, seq_len, hidden]
-
-        logits = _call_tagger(tagger, x, device)
-
-    pred_ids = logits.argmax(dim=-1)[0].tolist()
-    word_ids = _get_word_ids(tokenizer, words, tokenizer(words, is_split_into_words=True, return_tensors="pt",
-                                                        truncation=True, max_length=max_length), max_length)
-
-    results = []
-    seen = set()
-    for tok_i, w_i in enumerate(word_ids):
-        if w_i is None or w_i in seen:
-            continue
-        seen.add(w_i)
-        results.append((words[w_i], id2label[pred_ids[tok_i]]))
-
-    return results
-
-def find_label_vocab(vocabs):
-    for i, v in enumerate(vocabs):
-        if hasattr(v, "itos"):
-            itos = v.itos
-            if isinstance(itos, (list, tuple)) and any(x in itos for x in ["O", "B-PER", "I-PER"]):
-                return i, v
-    return None, None
-
-
-def _get_word_ids(tokenizer, words, enc, max_length):
-    # Fast tokenizers: BatchEncoding has word_ids()
-    if hasattr(enc, "word_ids"):
-        return enc.word_ids(batch_index=0)
-    # Fallback
-    return tokenizer(words, is_split_into_words=True, truncation=True, max_length=max_length).word_ids()
-
-def _call_tagger(tagger, x, device):
-    # Calls forward in a compatible way (x only vs x+labels, etc.)
-    params = list(inspect.signature(tagger.forward).parameters.keys())
-    # common: ['x']  or  ['x','labels',...]
-    if "labels" in params:
-        ignore_idx = getattr(tagger, "label_ignore_idx", 0)
-        labels = torch.full((x.size(0), x.size(1)), ignore_idx, dtype=torch.long, device=device)
-
-        kwargs = {}
-        if "segments_mask" in params:
-            kwargs["segments_mask"] = None
-        if "get_sent_repr" in params:
-            kwargs["get_sent_repr"] = False
-
-        out = tagger(x, labels, **kwargs)
-    else:
-        out = tagger(x)
-
-    # normalize outputs to logits tensor
-    if isinstance(out, (tuple, list)):
-        return out[-1]
-    if hasattr(out, "logits"):
-        return out.logits
-    return out
-
 label_vocab = label_vocab[0]  # the list loaded from pickle
 id2label = {i: s for i, s in enumerate(label_vocab.itos)}
 
-# idx, label_vocab = find_label_vocab(label_vocab)
-# print("label vocab index:", idx)
-# id2label = {i: s for i, s in enumerate(label_vocab.itos)}
-
 
 sentence = "ذهب احمد الى السوق"
-# id2label = {i: s for i, s in enumerate(label_vocab.itos)}
-# pairs = predict_ner(sentence, model, label_vocab, device="cpu")
-pairs = predict_ner(sentence, model, encoder, tokenizer, id2label, device="cpu")
-print(pairs)
+# Load tagger
+tagger, tag_vocab, train_config = load_checkpoint(checkpoint_path)
+
+# Convert text to a tagger dataset and index the tokens in args.text
+dataset, token_vocab = text2segments(sentence)
+
+vocabs = namedtuple("Vocab", ["tags", "tokens"])
+vocab = vocabs(tokens=token_vocab, tags=tag_vocab)
+
+# From the datasets generate the dataloaders
+dataloader = get_dataloaders(
+    (dataset,),
+    vocab,
+    args_path,
+    batch_size=32,
+    shuffle=(False,),
+)[0]
+
+# Perform inference on the text and get back the tagged segments
+segments = tagger.infer(dataloader)
+
+# Print results
+for segment in segments:
+    s = [
+        f"{token.text} ({'|'.join([t['tag'] for t in token.pred_tag])})"
+        for token in segment
+    ]
+    print(" ".join(s))