Add event argument extraction

main.py

@@ -599,6 +599,186 @@ def predict_re(request: RERequest):
     except Exception as e:
         return {"error": str(e)}
 
+
+# ============ Event Argument Extraction ==============
+from transformers import pipeline
+
+EVENT_MODEL_ID = "SinaLab/arabic-relation-extraction-model"
+EVENT_MAX_LEN = 128
+
+event_pipe = pipeline(
+    "sentiment-analysis",
+    model=EVENT_MODEL_ID,
+    tokenizer=EVENT_MODEL_ID,
+    device=0 if torch.cuda.is_available() else -1,
+    return_all_scores=True,
+    max_length=EVENT_MAX_LEN,
+    truncation=True
+)
+
+event_relation_prompt = {
+    "location": "مكان حدوث",
+    "agent": "أحد المتأثرين في",
+    "happened at": "تاريخ حدوث"
+}
+
+event_categories = {
+    "agent": ["PERS", "NORP", "OCC", "ORG"],
+    "location": ["LOC", "FAC", "GPE"],
+    "happened at": ["DATE", "TIME"]
+}
+
+event_relation_name_map = {
+    "agent": "hasAgent",
+    "location": "hasLocation",
+    "happened at": "hasDate"
+}
+
+
+def get_entity_category(entity_type, categories):
+    for category, types in categories.items():
+        if entity_type in types:
+            return category
+    return None
+
+
+def get_positive_score(predicted_relation):
+    """
+    The pipeline returns something like:
+    [
+      [
+        {"label": "LABEL_0", "score": 0.12},
+        {"label": "LABEL_1", "score": 0.88}
+      ]
+    ]
+
+    In your original code, you used:
+        predicted_relation[0][0]["score"]
+
+    If your positive class is LABEL_0, keep index 0.
+    If your positive class is LABEL_1, change this to index 1.
+
+    This version first tries LABEL_1, then falls back to index 0.
+    """
+
+    scores = predicted_relation[0]
+
+    for item in scores:
+        if item["label"] in ["LABEL_1", "relation", "RELATION", "positive"]:
+            return item["score"]
+
+    return scores[0]["score"]
+
+
+def event_argument_extractor(sentence):
+    entities = entities_and_types(sentence)
+
+    event_entities = [
+        (entity_name, entity_type)
+        for entity_name, entity_type in entities.items()
+        if entity_type == "EVENT"
+    ]
+
+    argument_entities = [
+        (entity_name, entity_type)
+        for entity_name, entity_type in entities.items()
+        if entity_type != "EVENT"
+    ]
+
+    output_list = []
+
+    for event_entity, event_type in event_entities:
+        for arg_name, arg_type in argument_entities:
+            category = get_entity_category(arg_type, event_categories)
+
+            if category not in event_relation_prompt:
+                continue
+
+            relation_sentence = (
+                f"[CLS] {sentence} [SEP] "
+                f"{event_entity} {event_relation_prompt[category]} {arg_name}"
+            )
+
+            predicted_relation = event_pipe(relation_sentence)
+            score = score = predicted_relation[0][0]["score"] #get_positive_score(predicted_relation)
+
+            if score > 0.0:
+                output_list.append({
+                    "Subject": {
+                        "Type": event_type,
+                        "Label": event_entity
+                    },
+                    "Relation": event_relation_name_map[category],
+                    "Object": {
+                        "Type": arg_type,
+                        "Label": arg_name
+                    },
+                    "Confidence": float(round(score, 4))
+                })
+
+    return output_list
+
+
+class EAERequest(BaseModel):
+    text: str
+
+
+@app.post("/predict_eae")
+def predict_eae(request: EAERequest):
+    try:
+        text = request.text.strip()
+
+        if not text:
+            return JSONResponse(
+                content={
+                    "resp": [],
+                    "statusText": "EMPTY_INPUT",
+                    "statusCode": 1,
+                },
+                media_type="application/json",
+                status_code=200,
+            )
+
+        sentences = sentence_tokenizer(
+            text,
+            dot=False,
+            new_line=True,
+            question_mark=False,
+            exclamation_mark=False
+        )
+
+        results = []
+
+        for sentence in sentences:
+            sentence = sentence.strip()
+            if not sentence:
+                continue
+
+            sentence_results = event_argument_extractor(sentence)
+            results.extend(sentence_results)
+
+        return JSONResponse(
+            content={
+                "resp": results,
+                "statusText": "OK",
+                "statusCode": 0,
+            },
+            media_type="application/json",
+            status_code=200,
+        )
+
+    except Exception as e:
+        return JSONResponse(
+            content={
+                "resp": [],
+                "statusText": str(e),
+                "statusCode": 500,
+            },
+            media_type="application/json",
+            status_code=500,
+        )
+
+
 # =========== Front End =============================
 from fastapi.staticfiles import StaticFiles
 from fastapi.responses import FileResponse