add json fallback to tgtalker

Author: ekmpa

examples/llm/TGTalker.py

@@ -17,7 +17,6 @@
 """
 
 import argparse
-import json
 import logging
 
 import numpy as np
@@ -31,6 +30,7 @@
     make_system_prompt,
     make_user_prompt,
     predict_link,
+    extract_destination_node,
 )
 from tqdm import tqdm
 from transformers import AutoModelForCausalLM, AutoTokenizer
@@ -207,7 +207,7 @@ def main() -> None:
 
                 try:
                     output = model(prompt, schema)
-                    pred_dst = int(json.loads(output)['destination_node'])
+                    pred_dst = extract_destination_node(output)
                     # Candidates: true destination followed by TGB negatives.
                     query_dst = torch.cat(
                         [batch.edge_dst[i].unsqueeze(0), batch.neg_batch_list[i]]

examples/llm/multihop.py

@@ -11,7 +11,6 @@
 """
 
 import argparse
-import json
 import logging
 
 import numpy as np
@@ -24,6 +23,7 @@
     make_multihop_user_prompt,
     make_system_prompt,
     predict_link,
+    extract_destination_node,
 )
 from tqdm import tqdm
 from transformers import AutoModelForCausalLM, AutoTokenizer
@@ -148,7 +148,7 @@ def main() -> None:
 
                 try:
                     output = model(prompt, schema)
-                    pred_dst = int(json.loads(output)['destination_node'])
+                    pred_dst = extract_destination_node(output)
                     query_dst = torch.cat(
                         [batch.edge_dst[i].unsqueeze(0), batch.neg_batch_list[i]]
                     )

examples/llm/tgtalker_utils.py

@@ -10,6 +10,8 @@
 from __future__ import annotations
 
 import collections
+import json
+import re
 from typing import Iterable, List, Sequence, Tuple
 
 import numpy as np
@@ -222,6 +224,38 @@ def predict_link(query_dst: torch.Tensor, llm_dst: int) -> torch.Tensor:
     return (query_dst == llm_dst).float()
 
 
+def extract_destination_node(output: object) -> int:
+    """Extract ``destination_node`` from model output.
+
+    Tries default strict JSON parse, with a minimal
+    fallback for slightly malformed text that still contains the key/value.
+    """
+    if output is None:
+        raise ValueError('Model output is None')
+
+    text = output if isinstance(output, str) else str(output)
+    text = text.strip()
+    if not text:
+        raise ValueError('Model output is empty')
+
+    try:
+        return int(json.loads(text)['destination_node'])
+    except Exception:
+        pass
+        
+    if isinstance(output, dict) and 'destination_node' in output:
+        return int(output['destination_node'])
+    if hasattr(output, 'destination_node'):
+        return int(getattr(output, 'destination_node'))
+
+    # Fallback for partially malformed JSON/text containing the key-value pair.
+    m = re.search(r'"destination_node"\s*:\s*(-?\d+)', text)
+    if m:
+        return int(m.group(1))
+
+    raise ValueError(f'Could not extract destination_node from output: {text[:160]}')
+
+
 class BackgroundBuffer:
     """Sliding window of the most recent global ``(src, dst, ts)`` edges.