add lang

generic_nel.py

@@ -145,21 +145,9 @@ class NelPipeline(Pipeline):
         :return:
         """
 
-        # {
-        #     "surface": sentences[i].split("[START]")[1].split("[END]")[0],
-        #     "lOffset": lOffset,
-        #     "rOffset": rOffset,
-        #     "type": "UNK",
-        #     "id": f"{lOffset}:{rOffset}:{surface}:{NEL_MODEL}",
-        #     "wkd_id": get_wikipedia_page_props(wikipedia_titles[i * 2]),
-        #     "wkpedia_pagename": wikipedia_titles[
-        #         i * 2
-        #         ],  # This can be improved with a real API call to get the QID
-        #     "confidence_nel": np.round(percentages[i], 2),
-        # }
         wikipedia_prediction, enclosed_entity, lOffset, rOffset, percentage = outputs
         qid, language = get_wikipedia_page_props(wikipedia_prediction)
-        title, url = get_wikipedia_title(qid, language="en")
+        title, url = get_wikipedia_title(qid, language=language)
 
         percentage = round(percentage, 2)
 
@@ -176,5 +164,4 @@ class NelPipeline(Pipeline):
                 "rOffset": rOffset,
             }
         ]
-        # print(results)
         return results