updated app.py with multiple models

app.py

@@ -1,45 +1,69 @@
 import gradio as gr
 import torch
-from transformers import AutoTokenizer, AutoModelForTokenClassification
-
-tokenizer = AutoTokenizer.from_pretrained("ai4bharat/IndicNER")
-
-model = AutoModelForTokenClassification.from_pretrained("ai4bharat/IndicNER")
-
-
-def get_ner(sentence):
-    tok_sentence = tokenizer(sentence, return_tensors='pt')
+from transformers import AutoTokenizer, AutoModelForTokenClassification,pipeline
 
+def get_ner_bio(pipe,text):
+    tok_text = pipe.tokenizer(text, return_tensors='pt')
     with torch.no_grad():
-        logits = model(**tok_sentence).logits.argmax(-1)
+        logits = pipe.model(**tok_text).logits.argmax(-1)
         predicted_tokens_classes = [
-            model.config.id2label[t.item()] for t in logits[0]]
+          pipe.model.config.id2label[t.item()] for t in logits[0]
+        ]
 
         predicted_labels = []
-
         previous_token_id = 0
-        word_ids = tok_sentence.word_ids()
+        word_ids = tok_text.word_ids()
         for word_index in range(len(word_ids)):
-            if word_ids[word_index] == None:
-                previous_token_id = word_ids[word_index]
-            elif word_ids[word_index] == previous_token_id:
-                previous_token_id = word_ids[word_index]
-            else:
-                predicted_labels.append(predicted_tokens_classes[word_index])
-                previous_token_id = word_ids[word_index]
-
-        ner_output = []
-        for index in range(len(sentence.split(' '))):
-            ner_output.append(
-                (sentence.split(' ')[index], predicted_labels[index]))
+          if not (word_ids[word_index] == None or word_ids[word_index] == previous_token_id):
+              predicted_labels.append(predicted_tokens_classes[word_index])
+          previous_token_id = word_ids[word_index]
+
+        ner_output = [
+            (word, label if label!="O" else None) 
+            for word, label in zip(text.split(" "),predicted_labels)
+        ]
         return ner_output
 
+def get_ner(pipe,text,aggregation_strategy="first"):
+
+  if aggregation_strategy == "bio_first":
+    return get_ner_bio(pipe,text)
+  else:
+    results = pipe(text,aggregation_strategy=aggregation_strategy)
+    for result in results:
+      result["entity"] = result["entity_group"]
+    return {"text": text, "entities": results} 
+
+ner_models = [
+  "livinNector/TryNER-500",
+  "livinNector/TryNER-1k",
+  "livinNector/IndicBERTNER",
+  "livinNector/IndicNER",
+  "ai4bharat/IndicNER",
+  "livinNector/distilbert-multilingual-base-ner"
+]
+ner_pipes = [pipeline("token-classification",model) for model in ner_models]
+
+def get_ner_outputs(text,aggregation_strategy):
+  return [get_ner(pipe,text,aggregation_strategy) for pipe in ner_pipes]
+examples = [
+        ["ஆனந்த் மற்றும் லிவின் நெக்டர் ஆகியொர் அண்ணாமலை பல்கலைக்கழகத்தில் படித்து வருகின்றனர்.","first"],
+        ["இந்தியன் இன்ஸ்டிட்யூட் ஆஃப் டெக்னாலஜி மெட்ராஸ் கிண்டியில் அமைந்துள்ளது.","average"],
+        ["சச்சின் டெண்டுல்கர் மும்பை மாநகரத்தைச் சேர்ந்த ஒரு நடுத்தரக் குடும்பத்தில் நான்காவது குழந்தையாகப் பிறந்தார். பல துடுப்பாட்ட வீரர்களை உருவாக்கிய சாரதாஷ்ரம் வித்யாமந்திர் பள்ளியில் சேர்ந்தார்.","bio_first"]
+
+    ]
 
-iface = gr.Interface(get_ner,
-                     gr.Textbox(placeholder="Enter sentence here..."),
-                     ["highlight"], description='NER Specialized for Tamil Language.',
-                     examples=["முதல்வர் ஸ்டாலின் பட்டமளிப்பு விழாவிற்காக சிதம்பரத்திலுள்ள அண்ணாமலைப் பல்கலைகழகத்திற்கு வருகை தந்தார்.","வல்லவராயன் வந்தியதேவனும் ஆதித்திய கரிகாலனும் கடம்பூருக்குச் சென்றனர். "], title='TaNER',
-                     article='TaNER is a model developed for NER in Tamil Language'
-                     )
+iface = gr.Interface(
+    get_ner_outputs,
+    [
+        gr.Textbox(value=examples[0][0]),
+        gr.Dropdown(["bio_first","first","max","average"],value=examples[0][1])
+     ],
+    [gr.Highlight(label=model) for model in ner_models], 
+    description='Named Entity Recongnition Interface Comparing Various Transformer Based NER models for Tamil Language.',
+    examples=examples, 
+    title='TaNER',
+    
+    )
 
-iface.launch(enable_queue=True)
+iface.launch(enable_queue=True)