added T5

app.py

@@ -1,5 +1,6 @@
 import gradio as gr
 
+from transformers import T5Tokenizer, T5ForConditionalGeneration
 from mailparser import parse_from_file
 from bs4 import BeautifulSoup
 from gliner import GLiNER
@@ -11,6 +12,9 @@ import os
 import en_core_web_sm
 nlp = en_core_web_sm.load()
 
+t5_tokenizer = T5Tokenizer.from_pretrained("t5-small")
+t5_model = T5ForConditionalGeneration.from_pretrained("t5-small")
+
 _MODEL = {}
 _CACHE_DIR = os.environ.get("CACHE_DIR", None)
 
@@ -58,6 +62,13 @@ def parse_query(sentences: List[str], labels: List[str], threshold: float = 0.3,
 
     return results
 
+def refine_entities_with_t5(entities):
+    inputs = "refine entities: " + " ; ".join([f"{entity['text']} as {entity['label']}" for entity in entities])
+    input_ids = t5_tokenizer.encode(inputs, return_tensors="pt", add_special_tokens=True)
+    outputs = t5_model.generate(input_ids)
+    result = t5_tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return result
+
 def present(email_file, labels, multilingual=False):
     email = accept_mail(email_file)
     cleaned_text = clean_email(email)
@@ -67,16 +78,18 @@ def present(email_file, labels, multilingual=False):
     entities = parse_query(sentence_list, labels, threshold=0.3, nested_ner=False, model_name="urchade/gliner_base", multilingual=multilingual)
     
     # Format entities for DataFrame: Convert list of dicts to list of lists
-    entities_data = [[entity['text'], entity['label']] for entity in entities]
+    entities = [[entity['text'], entity['label']] for entity in entities]
+
+    refined_entities = refine_entities_with_t5(entities)
 
     email_info = {
         "Subject": email.subject,
         "From": email.from_,
         "To": email.to,
         "Date": email.date,
-        "Extracted Entities": entities_data  # Adjusted for DataFrame
+        "Extracted Entities": refined_entities
     }
-    return [email_info[key] for key in ["Subject", "From", "To", "Date"]] + [entities_data]
+    return [email_info[key] for key in email_info]
 
 labels = ["PERSON", "PRODUCT", "DEAL", "ORDER", "ORDER PAYMENT METHOD", "STORE", "LEGAL ENTITY", "MERCHANT", "FINANCIAL TRANSACTION", "UNCATEGORIZED", "DATE"]
 
@@ -96,8 +109,10 @@ demo = gr.Interface(
         gr.components.Textbox(label="From"),
         gr.components.Textbox(label="To"),
         gr.components.Textbox(label="Date"),
-        gr.components.Dataframe(headers=["Text", "Label"], label="Extracted Entities")
+        gr.components.Dataframe(headers=["Text", "Label"], label="Extracted Entities"),
+        gr.components.Textbox(label="Refined Entities")
     ],
+    layout="horizontal",
     title="Email Info Extractor",
     description="Upload an email file (.eml) to extract its details and detected entities."
 )

requirements.txt

@@ -1,4 +1,5 @@
 gliner
+transformers
 mail-parser
 gradio
 beautifulsoup4