Views
No views yet
MMG/XLM-roberta-large-ner-spanish and was finetuned using boletines judiciales.1from transformers import AutoTokenizer, AutoModelForTokenClassification, pipeline
2
3REPO = "agomez302/nlp-dr-ner"
4
5class NerProcessor:
6 def __init__(self):
7 self.deployed_tokenizer = AutoTokenizer.from_pretrained(REPO)
8 self.deployed_model = AutoModelForTokenClassification.from_pretrained(REPO)
9 self.deployed_ner_pipeline = pipeline(
10 "ner",
11 model=self.deployed_model,
12 tokenizer=self.deployed_tokenizer,
13 aggregation_strategy="simple"
14 )
15
16 def process_text(self, text):
17 """Runs NER model on text and returns JSONL string."""
18 try:
19 chunks = self.split_text_with_overlap(text)
20 all_predictions = []
21 for chunk in chunks:
22 preds = self.deployed_ner_pipeline(chunk)
23 all_predictions.extend(preds)
24 all_predictions = self.deduplicate_entities(all_predictions)
25
26 formatted_output = {
27 "entities": self.run_predictions(all_predictions)
28 }
29
30 return json.dumps(formatted_output)
31
32 except Exception as e:
33 logger.error(f"Failed to run NER model on extracted text: {e}")
34
35 def split_text_with_overlap(self, text, max_tokens=450, overlap=50):
36 """Split text into chunks with overlap to handle long sequences."""
37 if not text:
38 return []
39 max_tokens = min(max_tokens, 512)
40
41 tokenizer = self.deployed_tokenizer
42 tokens = tokenizer.encode(text, truncation=False)
43
44 if len(tokens) <= max_tokens:
45 return [text]
46
47 chunks = []
48 i = 0
49 while i < len(tokens):
50 chunk = tokenizer.decode(tokens[i:i + max_tokens], skip_special_tokens=True)
51 chunks.append(chunk)
52 i += max_tokens - overlap
53 return chunks
54
55 def deduplicate_entities(self, predictions):
56 """Remove duplicate entities from overlapping chunks."""
57 unique = []
58 seen = set()
59 for entity in predictions:
60 key = (entity['entity_group'], entity['word'], entity['start'], entity['end'])
61 if key not in seen:
62 unique.append(entity)
63 seen.add(key)
64 return unique
65
66 def run_predictions(self, predictions: list):
67 """Format predictions for output, converting float32 to regular float."""
68 try:
69 processed_predictions = []
70 for pred in predictions:
71 pred_dict = dict(pred)
72 pred_dict['score'] = float(pred_dict['score'])
73 processed_predictions.append(pred_dict)
74
75 return processed_predictions
76
77 except Exception as e:
78 logging.error(f"Failed to process predictions: {e}")
79 raise
80
81
82def main():
83 text = "SENTENCIA DEL 31 DE ENERO DE 2024 ... que la sentencia que antecede fue dada y firmada por los jueces que figuran en ella, en la fecha arriba indicada. www.poderjudicial.gob.do\n"
84 ner_processor = NerProcessor()
85 ner_output = ner_processor.process_text(text)
86 print(ner_output)
87
88if __name__ = '__main__':
89 main()
901{
2 "entities":[
3 0:{
4 "entity_group":"DATE"
5 "score":0.9878288507461548
6 "word":"veintitrés (23) días del mes de mayo del año dos mil veintitrés (2023)"
7 "start":290
8 "end":360
9 }
10 1:{
11 "entity_group":"DATE"
12 "score":0.9994959831237793
13 "word":"23 de mayo del año 2023"
14 "start":1058
15 "end":1081
16 }
17}{"text": "SENTENCIA DEL 31 DE ENERO DE 2024 ... que la sentencia que antecede fue dada y firmada por los jueces que figuran en ella, en la fecha arriba indicada. www.poderjudicial.gob.do\n", "entities": [{"start": 113, "end": 132, "label": "DATE"}, {"start": 271, "end": 292, "label": "DATE"}, {"start": 2009, "end": 2029, "label": "DATE"}, {"start": 2246, "end": 2265, "label": "DATE"}, {"start": 3083, "end": 3102, "label": "DATE"}, {"start": 3281, "end": 3300, "label": "DATE"}, {"start": 3479, "end": 3497, "label": "DATE"}, {"start": 3569, "end": 3588, "label": "DATE"}, {"start": 3872, "end": 3891, "label": "DATE"}, {"start": 7936, "end": 7955, "label": "DATE"}]}
// and so forth with further json lines1 # Split the dataset into training and testing sets (e.g., 80% train, 20% test)
2 split_dataset = dataset.train_test_split(test_size=0.2)
3 train_dataset = split_dataset["train"]
4 validation_dataset = split_dataset["test"]