previous working version

app.py

@@ -1,6 +1,5 @@
 import gradio as gr
 
-from transformers import T5Tokenizer, T5ForConditionalGeneration
 from mailparser import parse_from_file
 from bs4 import BeautifulSoup
 from gliner import GLiNER
@@ -12,9 +11,6 @@ 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)
 
@@ -42,49 +38,57 @@ def get_sentences(further_cleaned_text):
 def get_model(model_name: str = None, multilingual: bool = False):
     if model_name is None:
         model_name = "urchade/gliner_base" if not multilingual else "urchade/gliner_multilingual"
-    if model_name not in _MODEL:
+
+    global _MODEL
+
+    if _MODEL.get(model_name) is None:
         _MODEL[model_name] = GLiNER.from_pretrained(model_name, cache_dir=_CACHE_DIR)
+
     return _MODEL[model_name]
 
-def parse_query(sentences, labels, threshold=0.3, nested_ner=False, model_name=None, multilingual=False):
-    model = get_model(model_name, multilingual)
+def parse_query(sentences: List[str], labels: List[str], threshold: float = 0.3, nested_ner: bool = False, model_name: str = None, multilingual: bool = False) -> List[Dict[str, Union[str, list]]]:
+    model = get_model(model_name, multilingual=multilingual)
+
     results = []
+
     for sentence in sentences:
         _entities = model.predict_entities(sentence, labels, threshold=threshold)
-        results.extend([{"text": entity["text"], "label": entity["label"]} for entity in _entities])
-    return results
+        entities = [{"text": entity["text"], "label": entity["label"]} for entity in _entities]
+        results.extend(entities)
 
-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
+    return results
 
 def present(email_file, labels, multilingual=False):
     email = accept_mail(email_file)
     cleaned_text = clean_email(email)
     further_cleaned_text = remove_special_characters(cleaned_text)
     sentence_list = get_sentences(further_cleaned_text)
+    
     entities = parse_query(sentence_list, labels, threshold=0.3, nested_ner=False, model_name="urchade/gliner_base", multilingual=multilingual)
-    refined_entities = refine_entities_with_t5(entities)
+    
+    # Format entities for DataFrame: Convert list of dicts to list of lists
+    entities_data = [[entity['text'], entity['label']] for entity in entities]
+
     email_info = {
         "Subject": email.subject,
         "From": email.from_,
         "To": email.to,
         "Date": email.date,
-        "Extracted Entities": entities,  # Prepare entities for DataFrame if needed
-        "Refined Entities": refined_entities
+        "Extracted Entities": entities_data  # Adjusted for DataFrame
     }
-    return [email_info[key] for key in ["Subject", "From", "To", "Date", "Extracted Entities", "Refined Entities"]]
+    return [email_info[key] for key in ["Subject", "From", "To", "Date"]] + [entities_data]
 
 labels = ["PERSON", "PRODUCT", "DEAL", "ORDER", "ORDER PAYMENT METHOD", "STORE", "LEGAL ENTITY", "MERCHANT", "FINANCIAL TRANSACTION", "UNCATEGORIZED", "DATE"]
 
 demo = gr.Interface(
-    fn=present,
+    fn=present, 
     inputs=[
         gr.components.File(label="Upload Email (.eml file)"),
-        gr.components.CheckboxGroup(choices=labels, label="Labels to Detect", value=labels),
+        gr.components.CheckboxGroup(
+            choices=labels,
+            label="Labels to Detect",
+            value=labels,  # Default all selected
+        ),
         gr.components.Checkbox(label="Use Multilingual Model")
     ],
     outputs=[
@@ -92,8 +96,7 @@ 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.Textbox(label="Refined Entities")
+        gr.components.Dataframe(headers=["Text", "Label"], label="Extracted Entities")
     ],
     title="Email Info Extractor",
     description="Upload an email file (.eml) to extract its details and detected entities."