Major fixes for huggingface endpoints

handler.py

@@ -2,10 +2,11 @@ from typing import Dict, List, Any
 from transformers import AutoConfig, AutoTokenizer
 from src.models import DNikudModel, ModelConfig
 from src.running_params import BATCH_SIZE, MAX_LENGTH_SEN
-from src.utiles_data import Nikud
-from src.models_utils import predict_single
+from src.utiles_data import Nikud, NikudDataset
+from src.models_utils import predict_single, predict
 import torch
 import os
+from tqdm import tqdm
 
 
 class EndpointHandler:
@@ -22,28 +23,90 @@ class EndpointHandler:
             len(Nikud.label_2_id["sin"]),
             device=self.DEVICE,
         ).to(self.DEVICE)
+        state_dict_model = self.model.state_dict()
+        state_dict_model.update(torch.load("./models/Dnikud_best_model.pth"))
+        self.model.load_state_dict(state_dict_model)
+        self.max_length = MAX_LENGTH_SEN
 
     def back_2_text(self, labels, text):
         nikud = Nikud()
         new_line = ""
+
         for indx_char, c in enumerate(text):
             new_line += (
                 c
-                + nikud.id_2_char(labels[0][1][1], "dagesh")
-                + nikud.id_2_char(labels[0][1][2], "sin")
-                + nikud.id_2_char(labels[0][1][0], "nikud")
+                + nikud.id_2_char(labels[indx_char][1][1], "dagesh")
+                + nikud.id_2_char(labels[indx_char][1][2], "sin")
+                + nikud.id_2_char(labels[indx_char][1][0], "nikud")
             )
             print(indx_char, c)
         print(labels)
         return new_line
 
+    def prepare_data(self, data, name="train"):
+        print("Data = ", data)
+        dataset = []
+        for index, (sentence, label) in tqdm(
+            enumerate(data), desc=f"Prepare data {name}"
+        ):
+            encoded_sequence = self.tokenizer.encode_plus(
+                sentence,
+                add_special_tokens=True,
+                max_length=self.max_length,
+                padding="max_length",
+                truncation=True,
+                return_attention_mask=True,
+                return_tensors="pt",
+            )
+            label_lists = [
+                [letter.nikud, letter.dagesh, letter.sin] for letter in label
+            ]
+            label = torch.tensor(
+                [
+                    [
+                        Nikud.PAD_OR_IRRELEVANT,
+                        Nikud.PAD_OR_IRRELEVANT,
+                        Nikud.PAD_OR_IRRELEVANT,
+                    ]
+                ]
+                + label_lists[: (self.max_length - 1)]
+                + [
+                    [
+                        Nikud.PAD_OR_IRRELEVANT,
+                        Nikud.PAD_OR_IRRELEVANT,
+                        Nikud.PAD_OR_IRRELEVANT,
+                    ]
+                    for i in range(self.max_length - len(label) - 1)
+                ]
+            )
+
+            dataset.append(
+                (
+                    encoded_sequence["input_ids"][0],
+                    encoded_sequence["attention_mask"][0],
+                    label,
+                )
+            )
+
+        self.prepered_data = dataset
+
     def predict_single_text(
         self,
         text,
     ):
-        data = self.tokenizer(text, return_tensors="pt")
-        all_labels = predict_single(self.model, data, self.DEVICE)
-        return all_labels
+        dataset = NikudDataset(tokenizer=self.tokenizer, max_length=MAX_LENGTH_SEN)
+        data, orig_data = dataset.read_single_text(text)
+        print("data", data, len(data))
+        dataset.prepare_data(name="inference")
+        mtb_prediction_dl = torch.utils.data.DataLoader(
+            dataset.prepered_data, batch_size=BATCH_SIZE
+        )
+        # print("dataset", dataset, len(dataset))
+        # data = self.tokenizer(text, return_tensors="pt")
+        all_labels = predict(self.model, mtb_prediction_dl, self.DEVICE)
+        text_data_with_labels = dataset.back_2_text(labels=all_labels)
+        # all_labels = predict_single(self.model, dataset, self.DEVICE)
+        return text_data_with_labels
 
     def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
         """
@@ -59,5 +122,5 @@ class EndpointHandler:
         # result = []
         # for pred in prediction:
         #     result.append(self.back_2_text(pred, inputs))
-        result = self.back_2_text(prediction, inputs)
-        return result
+        # result = self.back_2_text(prediction, inputs)
+        return prediction

models/Dnikud_best_model.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:31b1bb3dc66cebf70ad4bfa52d77257d92c745e6609f0023108e91041447e754
+size 446945642

src/models_utils.py

@@ -94,15 +94,17 @@ def predict_single(model, data, device="cpu"):
 
     all_labels = None
     with torch.no_grad():
-        inputs = data["input_ids"].to(device)
-        attention_mask = data["attention_mask"].to(device)
+        (inputs, attention_mask, labels_demo) = data
+        inputs = inputs.to(device)
+        attention_mask = attention_mask.to(device)
+        labels_demo = labels_demo.to(device)
 
-        # mask_cant_be_nikud = np.array(labels_demo.cpu())[:, :, 0] == -1
-        # mask_cant_be_dagesh = np.array(labels_demo.cpu())[:, :, 1] == -1
-        # mask_cant_be_sin = np.array(labels_demo.cpu())[:, :, 2] == -1
+        mask_cant_be_nikud = np.array(labels_demo.cpu())[:, :, 0] == -1
+        mask_cant_be_dagesh = np.array(labels_demo.cpu())[:, :, 1] == -1
+        mask_cant_be_sin = np.array(labels_demo.cpu())[:, :, 2] == -1
 
         nikud_probs, dagesh_probs, sin_probs = model(inputs, attention_mask)
-        print(nikud_probs, dagesh_probs, sin_probs)
+        print("model output: ", nikud_probs, dagesh_probs, sin_probs)
 
         pred_nikud = np.array(torch.max(nikud_probs, 2).indices.cpu()).reshape(
             inputs.shape[0], inputs.shape[1], 1
@@ -114,9 +116,9 @@ def predict_single(model, data, device="cpu"):
             inputs.shape[0], inputs.shape[1], 1
         )
 
-        # pred_nikud[mask_cant_be_nikud] = -1
-        # pred_dagesh[mask_cant_be_dagesh] = -1
-        # pred_sin[mask_cant_be_sin] = -1
+        pred_nikud[mask_cant_be_nikud] = -1
+        pred_dagesh[mask_cant_be_dagesh] = -1
+        pred_sin[mask_cant_be_sin] = -1
         # print(pred_nikud, pred_dagesh, pred_sin)
         pred_labels = np.concatenate((pred_nikud, pred_dagesh, pred_sin), axis=2)
         print(pred_labels)

src/utiles_data.py

@@ -370,6 +370,8 @@ class NikudDataset(Dataset):
         self.max_length = max_length
         self.tokenizer = tokenizer
         self.is_train = is_train
+        self.data = None
+        self.origin_data = None
         if folder is not None:
             self.data, self.origin_data = self.read_data_folder(folder, logger)
         elif file is not None:
@@ -453,6 +455,65 @@ class NikudDataset(Dataset):
 
         return data, orig_data
 
+    def read_single_text(self, text: str, logger=None) -> List[Tuple[str, list]]:
+        # msg = f"read file: {filepath}"
+        # if logger:
+        #     logger.debug(msg)
+        # else:
+        #     print(msg)
+        data = []
+        orig_data = []
+        # with open(filepath, "r", encoding="utf-8") as file:
+        #     file_data = file.read()
+        data_list = self.split_text(text)
+        # print("data_list", data_list)
+        for sen in tqdm(data_list, desc=f"Source: {data}"):
+            if sen == "":
+                continue
+
+            labels = []
+            text = ""
+            text_org = ""
+            index = 0
+            sentence_length = len(sen)
+            while index < sentence_length:
+                if (
+                    ord(sen[index]) == Nikud.nikud_dict["PUNCTUATION MAQAF"]
+                    or ord(sen[index]) == Nikud.nikud_dict["PUNCTUATION PASEQ"]
+                    or ord(sen[index]) == Nikud.nikud_dict["METEG"]
+                ):
+                    index += 1
+                    continue
+
+                label = []
+                l = Letter(sen[index])
+                if not (l.letter not in Nikud.all_nikud_chr):
+                    if sen[index - 1] == "\n":
+                        index += 1
+                        continue
+                assert l.letter not in Nikud.all_nikud_chr
+                if sen[index] in Letters.hebrew:
+                    index += 1
+                    while (
+                        index < sentence_length
+                        and ord(sen[index]) in Nikud.all_nikud_ord
+                    ):
+                        label.append(ord(sen[index]))
+                        index += 1
+                else:
+                    index += 1
+
+                l.get_label_letter(label)
+                text += l.normalized
+                text_org += l.letter
+                labels.append(l)
+
+            data.append((text, labels))
+            orig_data.append(text_org)
+        self.data = data
+        self.origin_data = orig_data
+        return data, orig_data
+
     def split_text(self, file_data):
         file_data = file_data.replace("\n", f"\n{unique_key}")
         data_list = file_data.split(unique_key)