Update app.py

app.py

@@ -1,13 +1,71 @@
+from transformers import AutoModelForTokenClassification, AutoTokenizer, pipeline
 import gradio as gr
 
-def ner(text):
-    api = gr.Interface.load("browndw/docusco-bert", src='models')
-    spans = api(text)
-    replaced_spans = [(key, None) if value=='O' else (key, value) for (key, value) in spans]
-    return replaced_spans
+def ner_tagging(text):
+  model_name = "browndw/docusco-bert"
+  tokenizer = AutoTokenizer.from_pretrained(model_name, add_prefix_space=True)
+  
+  model = AutoModelForTokenClassification.from_pretrained(model_name)
+  nlp = pipeline("ner", model=model, tokenizer=tokenizer)
+  ner_results = nlp(text.lower())
+  
+  output = []
+  
+  text_2 = text.split(" ")
+  
+  for i in range(len(text_2)):
+      ent = ner_results[i]["entity"]
+      if ent != "O":
+          output.extend([(text_2[i], ent), (" ", None)])
+      else:
+          output.extend([(text_2[i], None), (" ", None)])
 
-interface = gr.Interface(ner, inputs=gr.Textbox(label="Input", value="Autophagy maintains tumour growth through circulating arginine. Acute, \
-whole-body deletion of the essential autophagy gene Atg7 in adult mice causes a systemic metabolic defect that manifests as starvation  \
-intolerance and gradual loss of white adipose tissue, liver glycogen and muscle mass. Cancer cells also benefit from autophagy."), outputs="highlightedtext")
+  return output
 
-interface.launch()
+def get_entities(example):
+    model_name = "browndw/docusco-bert"
+    tokenizer = AutoTokenizer.from_pretrained(model_name, add_prefix_space=True)
+  
+    model = AutoModelForTokenClassification.from_pretrained(model_name)
+    token_classifier = pipeline("token-classification", aggregation_strategy="max", model=model, tokenizer=tokenizer)
+    results = token_classifier(example.lower())
+    
+    output = []
+
+    i=0
+    prev_item = None
+    next_item = None
+    while i < (len(results)):
+        item = results[i]
+        p=i-1
+        n=i+1
+        
+        if p > 0:
+            prev_item = results[p]
+        
+        
+        if n<(len(results)):
+            next_item = results[n]
+        
+        
+        if (i==0):
+            if item["start"]>0:
+                output.extend([(example[0:item["start"]], None)])
+        output.extend([(example[item["start"]:item["end"]], item["entity_group"])])
+        if (next_item!=None):
+            ##check the span
+            if(item["end"]!=next_item["start"]):
+                output.extend([(example[item["end"]:next_item["start"]], None)])
+        i=i+1
+
+    if item["end"] < len(example):
+        output.extend([(example[item["end"]:len(example)], None)])
+    
+    return output
+
+def greet(name):
+    return "Hello " + name + "!!"
+
+iface = gr.Interface(fn=get_entities, inputs="text", outputs=['highlight'], examples=[['Jaws is a splendidly shrewd cinematic equation which not only gives you one or two very nasty turns when you least expect them but, possibly more important, knows when to make you think another is coming without actually providing it.'], 
+['In order to understand how cyclic variations in turbulence intensities affect cycle-to-cycle variations in combustion, in-cylinder flow fields and turbulence need to be studied more closely.']], title='DocuScope Demo (BERT)', description = 'This is one of a family of models trained on DocuScope. Click on one of the examples below and SUBMIT. Be sure to CLEAR the output before tagging a new submission. You can also enter your own text.')
+iface.launch()