Update main.py

main.py

@@ -267,41 +267,18 @@ re_model = BertRE(num_labels=len(rel2id))
 re_model.load_state_dict(torch.load(weights_path, map_location="cpu"))
 re_model.eval()
 
-'''
-def entities_and_types(sentence):
-    ner_output = extract(sentence) 
-    entities = distill_entities(ner_output)
-
-    entity_dict = {}
-    for name, entity_type, _, _ in entities:
-        entity_dict[name] = entity_type
 
-    return entity_dict
-'''
 def convert_ner_format(ner_output):
     return [[item["token"], item["tags"]] for item in ner_output]
     
 def entities_and_types(sentence):
-    print("\n=== NER DEBUG ===")
-    print("INPUT:", sentence)
-
     ner_output = extract(sentence)
-    print("RAW NER OUTPUT:", ner_output)
-
-    # ✅ FIX HERE
     converted = convert_ner_format(ner_output)
-
-    print("CONVERTED FORMAT:", converted)
-
     entities = distill_entities(converted)
-
-    print("DISTILLED ENTITIES:", entities)
-
     entity_dict = {}
     for name, entity_type, _, _ in entities:
         entity_dict[name] = entity_type
 
-    print("ENTITY DICT:", entity_dict)
     return entity_dict
 
 relation_domain_range=[
@@ -514,8 +491,6 @@ for rel in relation_domain_range:
         for r in rel["range"]:
             relation_lookup[d][r].append(rel["relation"])
 
-
-'''
 def insert_markers(sentence, ent1, ent2):
     if ent1 not in sentence or ent2 not in sentence:
         return None
@@ -525,23 +500,7 @@ def insert_markers(sentence, ent1, ent2):
     marked = marked.replace(ent2, f"[Obj] {ent2} [/Obj]", 1)
 
     return marked
-'''
-def insert_markers(sentence, ent1, ent2):
-    print("\n--- MARKER DEBUG ---")
-    print("Original sentence:", sentence)
-    print("Entity1:", ent1)
-    print("Entity2:", ent2)
-
-    if ent1 not in sentence or ent2 not in sentence:
-        print("❌ Entity not found in sentence")
-        return None
 
-    marked = sentence
-    marked = marked.replace(ent1, f"[Sub] {ent1} [/Sub]", 1)
-    marked = marked.replace(ent2, f"[Obj] {ent2} [/Obj]", 1)
-
-    print("Marked sentence:", marked)
-    return marked
 def encode(sentence):
     enc = relation_tokenizer(
         sentence,
@@ -603,7 +562,18 @@ def relation_extractor(sentence):
             continue
 
         if conf > 0.80 and rel != "no_relation" and rel.split(".")[-1] in valid_rels:
-            output.append([ent1, rel, ent2, conf])
+            output.append({
+                "Subject": {
+                    "Type": type1,
+                    "Label": ent1
+                },
+                "Relation": rel,
+                "Object": {
+                    "Type": type2,
+                    "Label": ent2
+                },
+                "Confidence": float(round(conf, 4)) 
+            })
 
     return output