Update README.md

README.md

@@ -28,30 +28,29 @@ This work was funded by the Spanish Government, the Spanish Ministry of Economy
 from transformers import pipeline, AutoModelForTokenClassification, AutoTokenizer
 import torch
 
-def get_result_text_es_pt (list_entity, text, lang):
+def get_result_text_ca (list_entity, text):
     result_words = []
-    if lang == "es":
-        punc_tags = ['¿', '?', '¡', '!', ',', '.', ':']
-    else:
-        punc_tags = ['?', '!', ',', '.', ':']
+    punc_tags = ['?', '!', ',', '.', ':']
 
     for entity in list_entity:
-        tag = entity["entity"]
-        word = entity["word"]
         start = entity["start"]
         end = entity["end"]
+        tag = entity["entity"]
+        word = entity["word"]
 
         # check punctuation
         punc_in = next((p for p in punc_tags if p in tag), "")
 
         subword = False
         # check subwords
-        if word[0] == "#":
+        if word[0] != "Ġ":
             subword = True
             if punc_in != "":
                 word = result_words[-1].replace(punc_in, "") + text[start:end]
             else:
                 word = result_words[-1] + text[start:end]
+        else:
+            word = text[start:end]
 
         if tag == "l":
             word = word
@@ -71,18 +70,20 @@ def get_result_text_es_pt (list_entity, text, lang):
 
     return " ".join(result_words)
 
-lang = "es"
-model_path = "VOCALINLP/spanish_capitalization_punctuation_restoration_sanivert"
+
+lang = "ca"
+model_path = "VOCALINLP/catalan_capitalization_punctuation_restoration_sanivert"
 
 model = AutoModelForTokenClassification.from_pretrained(model_path)
 tokenizer = AutoTokenizer.from_pretrained(model_path)
 
 pipe = pipeline("token-classification", model=model, tokenizer=tokenizer)
-text = "el paciente presenta los siguientes síntomas náuseas vértigo disnea fiebre y dolor abdominal"
+
+text = "el pacient presenta els símptomes següents febre dispnea nàusees i vòmits"
 result = pipe(text)
 
 print("Source text: "+ text)
-result_text = get_result_text_es_pt(result, text, lang)
+result_text = get_result_text_ca(result, text)
 print("Restored text: " +result_text)
 ```