Update app.py

app.py

@@ -1,5 +1,6 @@
 
 import gradio as gr
+import pandas as pd
 from transformers import pipeline, AutoModelForQuestionAnswering, AutoTokenizer, AutoModelForTokenClassification
 
 # Load your custom model and tokenizer
@@ -11,19 +12,35 @@ qa_tokenizer = AutoTokenizer.from_pretrained(qa_model_name)
 
 token_classification_model = AutoModelForTokenClassification.from_pretrained(token_classification_model_name)
 token_classification_tokenizer = AutoTokenizer.from_pretrained(token_classification_model_name)
-
-
+def tabulazier(output):
+    output_comb = []
+    for ind, entity in enumerate(output):
+        if ind == 0:
+            output_comb.append(entity)
+        elif output[ind]["start"] == output[ind-1]["end"] and output[ind]["entity_group"] == output[ind-1]["entity_group"]:
+            output_comb[-1]["word"] = output_comb[-1]["word"] + output[ind]["word"]
+            output_comb[-1]["end"] = output[ind]["end"]
+        else:
+            output_comb.append(entity)
+
+    df = pd.DataFrame(output_comb)
+    df['word'] = df['word'].str.replace('#', '')
+    return df
+
+
+# Define a function for inference based on pipeline type
 def predict(pipeline_type, question, context):
     if pipeline_type == "question-answering":
         qa_pipeline = pipeline("question-answering", model=qa_model, tokenizer=qa_tokenizer)
         result = qa_pipeline(question=question, context=context)
-        response = [(result['answer'], result['score'])]
-        return response
+        response = [(result['answer'], result.get('score', None))]
+        return [response, response]
     elif pipeline_type == "token-classification":
-        token_classification_pipeline = pipeline("token-classification", model=token_classification_model, tokenizer=token_classification_tokenizer)
+        token_classification_pipeline = pipeline("ner", model=token_classification_model, tokenizer=token_classification_tokenizer, aggregation_strategy="simple")
         result = token_classification_pipeline(context)
-        highlighted_text = {"text": context, "entities": result} 
-        return gr.HighlightedText(highlighted_text)
+        highlighted_text = {"text": context, "entities": result}
+        table=tabulazier(result)
+        return [gr.HighlightedText(highlighted_text), table]
 
 # Create a Gradio Interface with dropdown and two text inputs
 iface = gr.Interface(
@@ -33,8 +50,11 @@ iface = gr.Interface(
         "text",
         "text"
     ],
-    outputs=gr.Highlight()
+    outputs=[gr.Highlight(), gr.Dataframe()]
 )
 
+
+              
+
 # Launch the interface
-iface.launch()
+iface.launch(debug=False)