t5 add

app.py

@@ -4,6 +4,7 @@ from mailparser import parse_from_file
 from bs4 import BeautifulSoup
 from gliner import GLiNER
 from typing import Dict, Union, List
+from transformers import T5Tokenizer, T5ForConditionalGeneration
 
 import spacy
 import re
@@ -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)
@@ -96,7 +107,8 @@ 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")
     ],
     title="Email Info Extractor",
     description="Upload an email file (.eml) to extract its details and detected entities."