t5 fix

app.py

@@ -42,24 +42,16 @@ 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"
-
-    global _MODEL
-
-    if _MODEL.get(model_name) is None:
+    if model_name not in _MODEL:
         _MODEL[model_name] = GLiNER.from_pretrained(model_name, cache_dir=_CACHE_DIR)
-
     return _MODEL[model_name]
 
-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)
-
+def parse_query(sentences, labels, threshold=0.3, nested_ner=False, model_name=None, multilingual=False):
+    model = get_model(model_name, multilingual)
     results = []
-
     for sentence in sentences:
         _entities = model.predict_entities(sentence, labels, threshold=threshold)
-        entities = [{"text": entity["text"], "label": entity["label"]} for entity in _entities]
-        results.extend(entities)
-
+        results.extend([{"text": entity["text"], "label": entity["label"]} for entity in _entities])
     return results
 
 def refine_entities_with_t5(entities):
@@ -74,34 +66,25 @@ def present(email_file, labels, multilingual=False):
     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)
-    
-    # Format entities for DataFrame: Convert list of dicts to list of lists
-    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": refined_entities
+        "Extracted Entities": entities,  # Prepare entities for DataFrame if needed
+        "Refined Entities": refined_entities
     }
-    return [email_info[key] for key in email_info]
+    return [email_info[key] for key in ["Subject", "From", "To", "Date", "Extracted Entities", "Refined Entities"]]
 
 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,  # Default all selected
-        ),
+        gr.components.CheckboxGroup(choices=labels, label="Labels to Detect", value=labels),
         gr.components.Checkbox(label="Use Multilingual Model")
     ],
     outputs=[
@@ -112,7 +95,6 @@ demo = gr.Interface(
         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."
 )