Spaces:
Sleeping
Sleeping
| #!/usr/bin/env python | |
| import os | |
| import re | |
| import string | |
| os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2" | |
| os.environ["TF_ENABLE_ONEDNN_OPTS"] = "0" | |
| from simpletransformers.ner import NERModel | |
| class BERTmodel: | |
| def __init__(self, normalization="full", wrds_per_pred=256): | |
| self.normalization = normalization | |
| self.wrds_per_pred = wrds_per_pred | |
| self.overlap_wrds = 32 | |
| self.valid_labels = ["O", "F", "C", "Q"] | |
| self.label_to_punct = {"F": "۔", "C": "،", "Q": "؟", "O": ""} | |
| self.model = NERModel( | |
| "bert", | |
| "/code/models/urdu", | |
| use_cuda=False, | |
| labels=self.valid_labels, | |
| args={"silent": True, "max_seq_length": 512}, | |
| ) | |
| self.patterns = { | |
| "partial": r"[ً-٠ٰ۟-ۤۧ-۪ۨ-ۭ،۔؟]+", | |
| "full": string.punctuation + "،؛؟۔٪ء‘’", | |
| } | |
| def punctuation_removal(self, text: str) -> str: | |
| if self.normalization == "partial": | |
| return re.sub(self.patterns[self.normalization], "", text).strip() | |
| else: | |
| return "".join(ch for ch in text if ch not in self.patterns[self.normalization]) | |
| def punctuate(self, text: str): | |
| text = self.punctuation_removal(text) | |
| splits = self.split_on_tokens(text) | |
| full_preds_lst = [self.predict(i["text"]) for i in splits] | |
| preds_lst = [i[0][0] for i in full_preds_lst] | |
| combined_preds = self.combine_results(text, preds_lst) | |
| punct_text = self.punctuate_texts(combined_preds) | |
| return punct_text | |
| def predict(self, input_slice): | |
| return self.model.predict([input_slice]) | |
| def split_on_tokens(self, text): | |
| wrds = text.replace("\n", " ").split() | |
| response = [] | |
| lst_chunk_idx = 0 | |
| i = 0 | |
| while True: | |
| wrds_len = wrds[i * self.wrds_per_pred : (i + 1) * self.wrds_per_pred] | |
| wrds_ovlp = wrds[ | |
| (i + 1) * self.wrds_per_pred : (i + 1) * self.wrds_per_pred + self.overlap_wrds | |
| ] | |
| wrds_split = wrds_len + wrds_ovlp | |
| if not wrds_split: | |
| break | |
| response_obj = { | |
| "text": " ".join(wrds_split), | |
| "start_idx": lst_chunk_idx, | |
| "end_idx": lst_chunk_idx + len(" ".join(wrds_len)), | |
| } | |
| response.append(response_obj) | |
| lst_chunk_idx += response_obj["end_idx"] + 1 | |
| i += 1 | |
| return response | |
| def combine_results(self, full_text: str, text_slices): | |
| split_full_text = full_text.replace("\n", " ").split(" ") | |
| split_full_text = [i for i in split_full_text if i] | |
| split_full_text_len = len(split_full_text) | |
| output_text = [] | |
| index = 0 | |
| if len(text_slices[-1]) <= 3 and len(text_slices) > 1: | |
| text_slices = text_slices[:-1] | |
| for slice in text_slices: | |
| slice_wrds = len(slice) | |
| for ix, wrd in enumerate(slice): | |
| if index == split_full_text_len: | |
| break | |
| if ( | |
| split_full_text[index] == str(list(wrd.keys())[0]) | |
| and ix <= slice_wrds - 3 | |
| and text_slices[-1] != slice | |
| ): | |
| index += 1 | |
| pred_item_tuple = list(wrd.items())[0] | |
| output_text.append(pred_item_tuple) | |
| elif ( | |
| split_full_text[index] == str(list(wrd.keys())[0]) and text_slices[-1] == slice | |
| ): | |
| index += 1 | |
| pred_item_tuple = list(wrd.items())[0] | |
| output_text.append(pred_item_tuple) | |
| assert [i[0] for i in output_text] == split_full_text | |
| return output_text | |
| def punctuate_texts(self, full_pred: list): | |
| punct_resp = [] | |
| for punct_wrd, label in full_pred: | |
| punct_wrd += self.label_to_punct[label] | |
| if punct_wrd.endswith("‘‘"): | |
| punct_wrd = punct_wrd[:-2] + self.label_to_punct[label] + "‘‘" | |
| punct_resp.append(punct_wrd) | |
| punct_resp = " ".join(punct_resp) | |
| if punct_resp[-1].isalnum(): | |
| punct_resp += "۔" | |
| return punct_resp | |
| class Urdu: | |
| def __init__(self): | |
| self.model = BERTmodel() | |
| def punctuate(self, data: str): | |
| return self.model.punctuate(data) | |