framework

Planning/data/finetune/combine_triplets.py

@@ -0,0 +1,93 @@
+import json
+import os
+
+
+path = "./prime/1000.json"
+with open(path, 'r') as f:
+    data = json.load(f)
+    
+def combine_triplets_by_boundary_no_loops(NTar):
+    combined = NTar[:]
+    changes = True  # Track if merges happen during an iteration
+
+    while changes:
+        changes = False  # Reset changes for this iteration
+        new_combined = []  # To hold the merged triplets
+        used = [False] * len(combined)  # Track which triplets have been processed
+
+        for i in range(len(combined)):
+            if used[i]:
+                continue  # Skip already processed triplets
+            
+            current_triplet = combined[i]  # Convert triplet to a mutable list
+            used[i] = True
+            merged = False  # Track if the current triplet is merged with another
+
+            for j in range(len(combined)):
+                if i != j and not used[j]:
+                    other_triplet = combined[j]
+
+                    # Check if the current triplet can be merged with another
+                    if current_triplet[-1] == other_triplet[0]:  # Current's last matches other's first
+                        current_triplet.extend(other_triplet[1:])  # Merge, excluding duplicate entity
+                        used[j] = True
+                        merged = True
+                        changes = True  # A merge occurred
+                        break
+                    elif current_triplet[0] == other_triplet[-1]:  # Current's first matches other's last
+                        current_triplet = other_triplet[:-1] + current_triplet  # Merge, excluding duplicate entity
+                        used[j] = True
+                        merged = True
+                        changes = True
+                        break
+
+            # After merging or if no merge happened, add the triplet to the new list
+            new_combined.append(current_triplet)
+
+        # Update the combined list with the newly merged triplets
+        combined = new_combined
+
+    return combined
+
+
+def combine_tar_ntar(tar, ntar):
+    for nt in ntar:
+        for i in range(len(tar)):
+            if tar[i][-1] == nt[0]:
+                tar[i] = tar[i]+nt[1:]
+            elif tar[i][0] == nt[-1]:
+                tar[i] = nt[:-1]+tar[i]
+    return tar
+
+def check_order(routes, target):
+    for i in range(len(routes)):
+        if routes[i][-1] != target:
+            if routes[i][0] != target:
+                raise ValueError(f"Wrong order: {routes[i]}")
+            else:
+                routes[i] = routes[i][::-1]
+    return routes
+
+routes_list = []
+restrictions_list = []
+for i in range(len(data)):
+    triplets = data[i]['Triplets']
+    target = data[i]['Target']
+    restrictions = data[i]['Restriction']
+    tar = []
+    ntar = []
+    for tp in triplets:
+        if target in tp:
+            tar.append(tp)
+        else:
+            ntar.append(tp)
+    if len(ntar) > 0:
+        ntar = combine_triplets_by_boundary_no_loops(ntar)
+        routes = combine_tar_ntar(tar, ntar)
+    else:
+        routes = tar
+    print(target)
+    routes = check_order(routes, target)
+    routes_list.append(routes)
+    restrictions_list.append(restrictions)
+

Planning/data/finetune/prime/1000.json

@@ -0,0 +1,55 @@
+[
+    {
+        "query": "Is the lens-specific intermediate filament-like protein, filensin, which is encoded by a gene expressed in the lens of camera-type eyes but not found in nasal cavity epithelial cells, involved in providing structural support to the ocular lens?",
+        "answer": {
+            "Triplets": [
+                [
+                    "anatomy",
+                    "expression present",
+                    "gene/protein"
+                ],
+                [
+                    "anatomy",
+                    "expression absent",
+                    "gene/protein"
+                ]
+            ],
+            "Restriction": {
+                "anatomy": [
+                    "lens of camera-type eyes",
+                    "nasal cavity epithelial cells"
+                ],
+                "gene/protein": [
+                    "filensin"
+                ],
+                "cellular_component": [
+                    "ocular lens"
+                ]
+            },
+            "Target": "gene/protein"
+        }
+    },
+    {
+        "query": "Which anatomical structures lack the expression of genes or proteins involved in the interaction with the fucose metabolism pathway?",
+        "answer": {
+            "Triplets": [
+                [
+                    "anatomy",
+                    "expression absent",
+                    "gene/protein"
+                ],
+                [
+                    "gene/protein",
+                    "interacts with",
+                    "pathway"
+                ]
+            ],
+            "Restriction": {
+                "pathway": [
+                    "fucose metabolism"
+                ]
+            },
+            "Target": "anatomy"
+        }
+    }
+]

Planning/data/get_train_data/get_llm_data.py

@@ -0,0 +1,237 @@
+"""
+Desc: This file is used to get the training data from the LLM
+
+"""
+import sys
+from pathlib import Path
+
+# Get the absolute path of the current script
+current_file = Path(__file__).resolve()
+project_root = current_file.parents[3]
+
+# Add the project root to the system path
+sys.path.append(str(project_root))
+
+from stark_qa import load_qa
+
+import argparse
+import os
+from openai import AzureOpenAI
+import json
+import openai
+from prompts import prompts
+
+
+
+
+
+
+"""
+
+MAG:
+sys_content: 478/query
+output: 45/query
+input: 25/query
+1000 queries
+
+total price:
+    1. o1: $13.29
+    2. o3mini: $0.97
+    3. deepseek-chat: $0.24
+    4. deepseek-reasoner: $0.49
+    
+Amazon:
+sys_content: 478/query
+
+"""
+
+# get the prompt for different datasets
+def get_sys_content(dataset_name):
+    """
+        input: 
+            dataset_name: the name of the dataset
+        output:
+            sys_content: the sys_content for the dataset
+    """
+    sys_content = prompts(dataset_name)
+    
+    
+    return sys_content
+
+# get the response from the llm
+def get_response(sys_content, user_content):
+    
+    messages = [{"role": "system", "content": sys_content},
+               {"role": "user", "content": user_content}
+               ]
+    
+    chat_completion = client.chat.completions.create(
+        messages=messages,
+        model=parameters['azure']['model'], # parameters['azure']['model'], parameters['openai']['model']
+        # temperature=0,
+        seed=576879897,
+    )
+    response = chat_completion.choices[0].message.content
+    
+    # print(messages)
+    # print(response)
+    
+    return response
+
+# save the outputs to json file
+def save_json(data, dataset_name):
+    """
+        input: 
+            data: the data to be saved
+            dataset_name: the name of the dataset
+    """
+    
+    file_dir = f"/home/yongjia/dgl/Yongjia/MOE/Reasoner/data/finetune/{dataset_name}"
+    os.makedirs(file_dir, exist_ok=True)
+    file_path = f"{file_dir}/1000_{parameters['azure']['model']}.json"
+    
+    with open(file_path, 'w') as f:
+        json.dump(data, f, indent=4)
+    print(f"Saved to {file_path}")
+
+# get the reasoning graphs for a dataset
+def get_rg(dataset_name):
+    """
+        input: 
+            dataset_name: the name of the dataset
+        output:
+            rg: the reasoning graph for the dataset
+    """
+    
+    # get the prompt for the dataset
+    sys_content = get_sys_content(dataset_name)
+    
+    # get qa dataset 
+    qa = load_qa(dataset_name)
+    train_qa = qa.get_subset('train')
+    
+    # we sample 1000 queries from the training set
+    pair_list = []
+    failure_count = 0
+    for i in range(1500):
+        query, q_id, ans_ids, _ = train_qa[i]
+    
+        # call the llm to get the reasoning graph
+        response = get_response(sys_content, query)
+        print(response)
+        
+        # process the response
+    
+        if dataset_name == 'prime':
+            output = {
+                "Triplets":[],
+                "Restriction": [],
+                "Target": ""
+            }
+            
+            try:
+                response = response.split('\n')
+                triplets_raw = response[0].replace('Triplets:', '').strip()
+                triplets = json.loads(triplets_raw)
+                output['Triplets'] = triplets
+                
+                restriction_raw = response[1].replace('Restriction:', '').strip()
+                restriction = json.loads(restriction_raw)
+                output['Restriction'] = restriction
+                
+                target = response[2].replace('Target:', '').strip()
+                output['Target'] = target
+            except:
+                failure_count += 1
+                continue
+            
+        elif dataset_name == 'mag' or dataset_name == 'amazon':
+            output = {
+                "Metapath": "",
+                "Restriction": [],
+            }
+            
+            try:
+                response = response.split('\n')
+                metapath = response[0].replace('Metapath:', '').strip()
+                output['Metapath'] = metapath
+                
+                restriction_raw = response[1].replace('Restriction:', '').strip()
+                restriction = json.loads(restriction_raw)
+                output['Restriction'] = restriction
+            except:
+                failure_count += 1
+                continue
+            
+        else:
+            raise ValueError('The dataset is not supported')
+        
+        pair = {'query': query, 'answer': output}
+
+        pair_list.append(pair)
+        
+        if len(pair_list) == 1000:
+            break
+            
+    # save the output to json file
+    save_json(pair_list, dataset_name)
+    print(f"Failure count: {failure_count}")
+    
+
+if __name__ == '__main__':    
+    # Argument parser setup
+    parser = argparse.ArgumentParser(description="Load LLM parameters and initialize API clients.")
+    
+    # Dataset name
+    parser.add_argument("--dataset_name", type=str, required=True, 
+                        choices=["mag", "amazon", "prime"], 
+                        help="Specify the dataset to use.")
+
+    # Model selection
+    parser.add_argument("--model", type=str, required=True, 
+                        choices=["gpt-4o-mini-20240718", "gpt-4o-2024-05-13", 
+                                "deepseek-reasoner", "gpt-o1-2024-12-17", 
+                                "o3-mini-2025-01-31"],
+                        help="Specify the model to use.")
+
+    # Azure API parameters
+    parser.add_argument("--azure_api_key", type=str, default=None, help="Azure API Key")
+    parser.add_argument("--azure_endpoint", type=str, default=None, help="Azure API Endpoint")
+    parser.add_argument("--azure_api_version", type=str, default=None, help="Azure API Version")
+
+    # OpenAI API parameters
+    parser.add_argument("--openai_api_key", type=str, default=None, help="OpenAI API Key")
+    parser.add_argument("--openai_endpoint", type=str, default=None, help="OpenAI API Endpoint")
+
+    args = parser.parse_args()
+
+    # Initialize parameters dictionary
+    parameters = {
+        "azure": {
+            "api_key": args.azure_api_key,
+            "azure_endpoint": args.azure_endpoint,
+            "api_version": args.azure_api_version,
+        },
+        "openai": {
+            "api_key": args.openai_api_key,
+            "endpoint": args.openai_endpoint,
+        }
+    }
+
+
+    # Determine which API client to use
+    if parameters["openai"]["api_key"]:
+        client = openai.OpenAI(
+            base_url=parameters["openai"]["endpoint"],
+            api_key=parameters["openai"]["api_key"],
+        )
+    else:
+        client = AzureOpenAI(
+            azure_endpoint=parameters["azure"]["azure_endpoint"],
+            api_key=parameters["azure"]["api_key"],
+            api_version=parameters["azure"]["api_version"],
+        )
+    
+    get_rg(args.dataset_name)
+
+

Planning/data/get_train_data/post_process_data.py

@@ -0,0 +1,58 @@
+"""
+Description: This script is used to post-process the data after getting it from the LLM.
+
+input: data from llm
+output: data for llama finetuning
+
+"""
+import json
+import os
+import copy
+
+def save_data(data, dataset_name, model_name):
+    file_dir = f"../data/finetune/{dataset_name}"
+    os.makedirs(file_dir, exist_ok=True)
+    file_path = os.path.join(file_dir, f"llama_ft_{model_name}.jsonl")
+    
+    with open(file_path, "w") as f:
+        for d in data:
+            f.write(json.dumps(d) + "\n")
+    
+    print(f"Data saved to {file_path}")
+    
+def process(sample, dataset_name, model_name):
+    """
+    input: sample from llm
+    output: sample for llama finetuning
+    """
+   
+    output_format = {"conversations": [{"role": "user", "content": ""}, {"role": "assistant", "content": ""}]}
+    ft_list = []
+    for i in range(len(sample)):
+        output = copy.deepcopy(output_format)
+        output["conversations"][0]["content"] = sample[i]["query"]
+        output["conversations"][1]["content"] = str(sample[i]["answer"])        
+        ft_list.append(output)
+    
+    # save data
+    save_data(ft_list, dataset_name, model_name)
+        
+    
+# ***** Main *****
+if __name__ == "__main__":
+    # read data
+    dataset_name_list = ["mag"]
+    model_names = ["gpt-4o-mini-20240718", "o3-mini-2025-01-31", "gpt-o1-2024-12-17", "gpt-4o-2024-05-13", "gpt-4o-mini-20240718", "gpt35-1106"] # gpt-o1-2024-12-17, "gpt-4o-mini-20240718", "gpt35-1106", o3-mini-2025-01-31
+    
+    for dataset_name in dataset_name_list:
+        for model_name in model_names:
+            relative_path = f"finetune/{dataset_name}/1000.json" #  f"finetune/{dataset_name}/1000_{dataset_name}.json"
+            current_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))  # Directory of the current script
+            file_path = os.path.join(current_dir, relative_path)
+            with open(f"{file_path}", "r") as f:
+                sample = json.load(f)
+            
+            # process data
+            process(sample, dataset_name, model_name)
+            print(f"Processing {model_name} for {dataset_name} is done.")
+            break

Planning/data/get_train_data/prompts.py

@@ -0,0 +1,274 @@
+def prompts(dataset_name):
+    
+    """
+    
+    input: 
+        dataset_name: the name of the dataset
+    output:
+        prompt: the prompt for the dataset
+    """
+    if dataset_name == "amazon":
+        sys_prompt = """
+       
+        You are a reasoning graph finder agent. Your role is to:
+        1. Identify the underlying **meta-path** from a given question, which consists of the **entity types** at each reasoning step. 
+        2. Extract the **content restriction** for each **entity type** based on the question. If there is no restriction for an entity type, leave its value empty.
+
+        You will be provided with a predefined **Entity Type List**. Only use the entity types from this list when constructing the meta-path and restrictions. Your response must be concise and strictly adhere to the specified **Output Format**.
+        
+        """
+
+        # Define the entity type list
+        entity_type_list = """
+        Entity Type List:
+        - **brand**
+        - **category**
+        - **product**
+        
+        """
+        
+        demonstrations = """
+        Here are some examples:
+        **Question 1:**  
+        What are some 5x4 inch team sports decals by Football Fanatics that are easy to apply on exterior surfaces?  
+
+        **Answer 1:**  
+        Metapath: brand -> product  
+        Restriction: {"brand": ["Football Fanatics"], "product": ["5x4 inch team sports decals by Football Fanatics that are easy to apply on exterior surfaces"]}
+
+        **Question 2:**  
+        Looking for men's insulated pants suitable for heavy rain, up to 10k water resistance, and with reinforced cuffs specifically for added protection in parking lots. Any suggestions?  
+
+        **Answer 2:**  
+        Metapath: category -> product  
+        Restriction: {"category": ["pants"], "product": ["men\'s insulated pants suitable for heavy rain", "reinforced cuffs specifically for added protection in parking lots"]}
+
+        **Question 3:**  
+        Can you recommend a dive mask that would work well with the Dive Mask Freediving Mask Spearfishing Mask Low Volume Mini Mask? I need the two masks to be compatible for my diving activities.  
+
+        **Answer 3:**  
+        Metapath: product -> product  
+        Restriction: {"product": ["work well with the Dive Mask Freediving Mask Spearfishing Mask Low Volume Mini Mask", "compatible for diving activities", "a dive mask"]}
+
+        **Question 4:**  
+        What are some UV protective women's golf jackets from PUMA? I'm cautious about skin protection, but I'm a big fan of PUMA.  
+
+        **Answer 4:**  
+        Metapath: brand -> product <- category  
+        Restriction: {"brand": ["PUMA"], "category": ["golf jackets"], "product": ["UV protective women\'s golf jackets from PUMA", "skin protection"]}
+        
+        
+        
+        """
+        
+        output_format = """
+        
+        **Output Format**
+        Metapath: "",
+        Restriction: {}
+        """
+        
+        sys_content = sys_prompt + entity_type_list + output_format + demonstrations
+        
+        
+    elif dataset_name == "mag":
+        
+        sys_prompt = """
+        You are a reasoning finder agent. Your role is to:
+        1. Identify the underlying **meta-path** from a given question, which consists of the **entity types** at each reasoning step. 
+        2. Extract the **content restriction** for each **entity type** based on the question. If there is no restriction for an entity type, leave its value empty.
+
+        You will be provided with a predefined **Entity Type List**. Only use the entity types from this list when constructing the meta-path and restrictions. Your response must be concise and strictly adhere to the specified **Output Format**.
+        
+        """
+
+        entity_type_list = """
+        Entity Type List:
+        - **paper**
+        - **author**
+        - **institution**
+        - **field_of_study**
+        
+        """
+        
+        output_format = """
+        **Output Format**
+        Metapath: "",
+        Restriction: {}
+        """
+        
+        
+        demonstrations = """
+        
+        Here are some examples:
+        **Question 1:**  
+        Show me research articles on the association of quasi-periodic oscillations (QPOs) with noise in celestial bodies within the context of Bicoherence.
+
+        **Answer 1:**  
+        Metapath: field_of_study -> paper  
+        Restriction: {"field_of_study": ["oscillations", "physics"], "paper": ["association of quasi-periodic oscillations (QPOs) with noise in celestial bodies within the context of Bicoherence"]}
+
+        **Question 2:**  
+        What research on water absorption in different frequency ranges have been referenced or deemed significant in the paper entitled 'High-resolution terahertz atmospheric water vapor continuum measurements?
+
+        **Answer 2:**  
+        Metapath: paper -> paper
+        Restriction: {"paper": ["water absorption in different frequency ranges", "High-resolution terahertz atmospheric water vapor continuum measurements"]}
+
+        **Question 3:**  
+        Show me publications by A.J. Turvey on the topic of supersymmetry particle searches.
+
+        **Answer 3:**  
+        Metapath: author -> paper
+        Restriction: {"author": ["A.J. Turvey"], "paper": ["supersymmetry particle searches"]}
+
+        **Question 4:**  
+        Looking for papers co-authored by someone involved in "Strain transferring mechanism analysis of the substrate-bonded FBG sensor". The papers should be in the similar field which is fiber optic strain sensors, and further discuss their development and application.
+
+        **Answer 4:**  
+        Metapath: paper -> author -> paper
+        Restriction: {"paper": ["Strain transferring mechanism analysis of the substrate-bonded FBG sensor", "in the similar field which is fiber optic strain sensors, and further discuss their development and application"]}
+
+        **Question 5:**
+        Can you find any publications by the authors of "Tradeoffs in the Realization of Electrically Pumped Vertical External Cavity Surface Emitting Lasers," that delve into the topic of hybrid quantum well/quantum dot structures in the context of lasers?
+
+        **Answer 5:**
+        Metapath: paper -> author -> paper <- field_of_study
+        Restriction: {"paper": ["Tradeoffs in the Realization of Electrically Pumped Vertical External Cavity Surface Emitting Lasers", "hybrid quantum well/quantum dot structures in the context of lasers"], "field_of_study": ["lasers", "hybrid quantum well/quantum dot", "physics", "Emission spectrum"]}
+
+        **Question 6:**
+        Which publications from Altair Engineering authors focus on improving directional sensitivity across a wide range of frequencies?
+
+        **Answer 6:**
+        Metapath: institution -> author -> paper
+        Restriction: {"institution": ["Altair Engineering"], "paper": ["improving directional sensitivity across a wide range of frequencies"]}
+
+        **Question 7:**
+        Publications by Carlisle Companies authors on satellite instrumentation and space performance
+
+        **Answer 7:**
+        Metapath: institution -> author -> paper <- field_of_study
+        Restriction: {"institution": ["Carlisle Companies"], "paper": ["satellite instrumentation and space performance"], "field_of_study": ["satellite", 'space performance']}
+
+        """
+        
+        sys_content = sys_prompt + entity_type_list + output_format + demonstrations
+        
+    elif dataset_name == "prime":
+        
+        sys_prompt = """
+        You are a triplets extractor. Given a list of triplets and a query, please 
+        1. extract the triplets contained in the query and give a list to me. 
+        2. make a restriction list which contains the description of the entity in the query.
+        3. tell me which entity the query is asking for.
+        Your response must be concise and strictly adhere to the specified **Output Format**.
+        
+        """
+
+        triplets = """
+        Triplets list:
+        [('anatomy', 'expression absent', 'gene/protein'),
+        ('anatomy', 'expression present', 'gene/protein'),
+        ('anatomy', 'parent-child', 'anatomy'),
+        ('biological_process', 'interacts with', 'exposure'),
+        ('biological_process', 'interacts with', 'gene/protein'),
+        ('biological_process', 'parent-child', 'biological_process'),
+        ('cellular_component', 'interacts with', 'exposure'),
+        ('cellular_component', 'interacts with', 'gene/protein'),
+        ('cellular_component', 'parent-child', 'cellular_component'),
+        ('disease', 'associated with', 'gene/protein'),
+        ('disease', 'contraindication', 'drug'),
+        ('disease', 'indication', 'drug'),
+        ('disease', 'linked to', 'exposure'),
+        ('disease', 'off-label use', 'drug'),
+        ('disease', 'parent-child', 'disease'),
+        ('disease', 'phenotype absent', 'effect/phenotype'),
+        ('disease', 'phenotype present', 'effect/phenotype'),
+        ('drug', 'carrier', 'gene/protein'),
+        ('drug', 'contraindication', 'disease'),
+        ('drug', 'enzyme', 'gene/protein'),
+        ('drug', 'indication', 'disease'),
+        ('drug', 'off-label use', 'disease'),
+        ('drug', 'side effect', 'effect/phenotype'),
+        ('drug', 'synergistic interaction', 'drug'),
+        ('drug', 'target', 'gene/protein'),
+        ('drug', 'transporter', 'gene/protein'),
+        ('effect/phenotype', 'associated with', 'gene/protein'),
+        ('effect/phenotype', 'parent-child', 'effect/phenotype'),
+        ('effect/phenotype', 'phenotype absent', 'disease'),
+        ('effect/phenotype', 'phenotype present', 'disease'),
+        ('effect/phenotype', 'side effect', 'drug'),
+        ('exposure', 'interacts with', 'biological_process'),
+        ('exposure', 'interacts with', 'cellular_component'),
+        ('exposure', 'interacts with', 'gene/protein'),
+        ('exposure', 'interacts with', 'molecular_function'),
+        ('exposure', 'linked to', 'disease'),
+        ('exposure', 'parent-child', 'exposure'),
+        ('gene/protein', 'associated with', 'disease'),
+        ('gene/protein', 'associated with', 'effect/phenotype'),
+        ('gene/protein', 'carrier', 'drug'),
+        ('gene/protein', 'enzyme', 'drug'),
+        ('gene/protein', 'expression absent', 'anatomy'),
+        ('gene/protein', 'expression present', 'anatomy'),
+        ('gene/protein', 'interacts with', 'biological_process'),
+        ('gene/protein', 'interacts with', 'cellular_component'),
+        ('gene/protein', 'interacts with', 'exposure'),
+        ('gene/protein', 'interacts with', 'molecular_function'),
+        ('gene/protein', 'interacts with', 'pathway'),
+        ('gene/protein', 'ppi', 'gene/protein'),
+        ('gene/protein', 'target', 'drug'),
+        ('gene/protein', 'transporter', 'drug'),
+        ('molecular_function', 'interacts with', 'exposure'),
+        ('molecular_function', 'interacts with', 'gene/protein'),
+        ('molecular_function', 'parent-child', 'molecular_function'),
+        ('pathway', 'interacts with', 'gene/protein'),
+        ('pathway', 'parent-child', 'pathway')]
+         
+        """
+        
+        output_format = """
+        **Output Format**
+        Triplets: []
+        Restriction: {}
+        Target: ""
+        
+        """
+
+        demonstrations = """
+        Here are some examples:
+        **Question 1:**
+        Search for conditions that lack any associated treatment medications and have a connection to the formation of Onion bulb structures.
+        
+        **Answer 1:**
+        Triplets: [["effect/phenotype", "phenotype present", "disease"], ["drug", "off-label use", "disease"]]
+        Restriction: {'effect/phenotype': ['formation of Onion bulb structures']}
+        Target: disease
+
+        **Question 2:**
+        What drug should be avoided for focal hand dystonia and also targets the ORM2 gene/protein?
+        
+        **Answer 2:**
+        Triplets: [["disease", "contraindication", "drug"], ["drug", "target", "gene/protein"]]
+        Restriction: {"disease": ["focal hand dystonia"], "gene/protein": ["ORM2"]}
+        Target: drug
+
+        **Question 3:**
+        Could my hip muscle weakness be a sign of the mitochondrial disease my mother has, or another related condition?
+        Triplets: [["disease", "parent-child", "disease"], ["effect/phenotype", "phenotype present", "disease"]]
+        
+        **Answer 3:**
+        Restriction: {"anatomy": ["hip muscle"], "disease": ["mitochondrial disease my mother has", "another related condition"], "effect/phenotype": ["weakness"]}
+        Target: disease
+
+        **Question 4:**
+        Can you find the genes and proteins involved with the activity of dolichyl-phosphate-mannose-dependent alpha-1,6-mannosyltransferase?
+        
+        **Answer 4:**
+        Triplets: [["molecular_function", "interacts with", "gene/protein"]]
+        Restriction: {"molecular_function": ["dolichyl-phosphate-mannose-dependent alpha-1,6-mannosyltransferase"]}
+        Target: gene/protein
+        """
+        
+        sys_content = sys_prompt + triplets + output_format + demonstrations
+    
+    return sys_content

Planning/data/train_eval.py

@@ -0,0 +1,223 @@
+"""
+Description: This file is used to train and evaluate the llama model.
+
+"""
+import os
+from unsloth import FastLanguageModel
+from trl import SFTTrainer
+from transformers import TrainingArguments, DataCollatorForSeq2Seq
+from unsloth import is_bfloat16_supported
+from datasets import load_dataset
+from unsloth.chat_templates import train_on_responses_only
+from unsloth.chat_templates import get_chat_template
+from sklearn.model_selection import train_test_split
+from transformers import TextStreamer
+import ast
+import contractions
+import re
+from utils import remove_inner_single_quotes
+
+# *****load model and tokenizer*****
+max_seq_length = 2048
+dtype = None    # None for auto detection. Float16 for Tesla T4, V100, Bfloat16 for Ampere+
+load_in_4bit = True # Use 4bit quantization to reduce memory usage.
+
+model, tokenizer = FastLanguageModel.from_pretrained(
+    model_name = "unsloth/Llama-3.2-3B-Instruct",
+    max_seq_length = max_seq_length,    # specify the maximum length of input the model can accept
+    dtype = dtype,
+    load_in_4bit = load_in_4bit,
+)
+
+# add LoRA adapters so we only need to update 1 to 10% of all parameters!
+model = FastLanguageModel.get_peft_model(
+    model,
+    r = 16, # Choose any number > 0 ! Suggested 8, 16, 32, 64, 128
+    target_modules = ["q_proj", "k_proj", "v_proj", "o_proj",# Reference: https://colab.research.google.com/drive/1Ys44kVvmeZtnICzWz0xgpRnrIOjZAuxp?usp=sharing#scrollTo=2eSvM9zX_2d3
+                      "gate_proj", "up_proj", "down_proj",],
+    lora_alpha = 16,
+    lora_dropout = 0, # Supports any, but = 0 is optimized
+    bias = "none",    # Supports any, but = "none" is optimized
+    # [NEW] "unsloth" uses 30% less VRAM, fits 2x larger batch sizes!
+    use_gradient_checkpointing = "unsloth", # True or "unsloth" for very long context
+    random_state = 3407,
+    use_rslora = False,  # We support rank stabilized LoRA
+    loftq_config = None, # And LoftQ
+)
+
+tokenizer = get_chat_template(
+    tokenizer,
+    chat_template = "llama-3.1",
+)
+
+def formatting_prompts_func(data):
+    convos = data['conversations']
+    texts = [tokenizer.apply_chat_template(convo, tokenize = False, add_generation_prompt = False) for convo in convos]
+    
+    return { "text" : texts, }
+
+def train(train_dataset, val_dataset):
+    # format the dataset
+    texts_train = train_dataset.map(formatting_prompts_func, batched=True)
+    print(texts_train)
+    texts_val = val_dataset.map(formatting_prompts_func, batched=True)
+    print(texts_val)
+    
+    # load trainer
+    trainer = SFTTrainer(
+        model = model,
+        tokenizer = tokenizer,
+        train_dataset = texts_train,
+        eval_dataset= texts_val,
+        dataset_text_field = "text",
+        max_seq_length = max_seq_length,
+        data_collator=DataCollatorForSeq2Seq(tokenizer=tokenizer),  # A function or mechanism to batch and prepare data during training and evaluation
+        dataset_num_proc = 2,   # number of processes to use for data preprocessing
+        packing = False, # Can make training 5x faster for short sequences.
+        args = TrainingArguments(
+            per_device_train_batch_size = 4, 
+            gradient_accumulation_steps = 8,
+            warmup_steps = 5,
+            num_train_epochs = 100, # Set this for 1 full training run. # TODO: increase the value
+            max_steps = 1000, # TODO: increase the value
+            learning_rate = 2e-4,
+            fp16 = not is_bfloat16_supported(),
+            bf16 = is_bfloat16_supported(),
+            logging_steps = 1,
+            optim = "adamw_8bit",
+            weight_decay = 0.01,
+            lr_scheduler_type = "linear",
+            seed = 3407,
+            output_dir = "outputs",
+            do_eval=True,
+            report_to='wandb',
+            evaluation_strategy="epoch",    # Specifies that evaluations will happen at the end of each epoch, using texts_val for metrics calculation. If we set the evaluation strategy without passing evaluation dataset, there is an error.
+            save_strategy="epoch",      # Save checkpoints based on epoch.
+            load_best_model_at_end=True,    # Load the best one from the disk to memory. Otherwise, the model is still the one trained after last epoch.
+            metric_for_best_model="loss",       # Evaluation metric on evaluation set. Make sure the metric is an option included in TrainingArguments.
+            greater_is_better=False,     # The metric is "greater_is_better".
+            save_total_limit=1      # How many checkpoints will be saved.
+        ),
+    )
+    
+    # train only on outputs and ignore the loss of user's inputs
+    trainer = train_on_responses_only(
+        trainer,
+        instruction_part = "<|start_header_id|>user<|end_header_id|>\n\n",
+        response_part = "<|start_header_id|>assistant<|end_header_id|>\n\n",
+    )
+
+    trainer_stats = trainer.train()
+    
+    # save the model
+    checkpoint_dir = f"./checkpoints/{dataset_name}"
+    os.makedirs(checkpoint_dir, exist_ok=True)
+    checkpoint_path = os.path.join(checkpoint_dir, f"lora_model")
+    model.save_pretrained(checkpoint_path) # Local saving 
+    tokenizer.save_pretrained(checkpoint_path)
+    
+    print("Training completed.")
+    
+    return checkpoint_path
+
+def evaluate(test_dataset, checkpoint_path):
+    
+    # load model
+    model, tokenizer = FastLanguageModel.from_pretrained(
+        model_name = checkpoint_path,
+        max_seq_length = max_seq_length,    # specify the maximum length of input the model can accept
+        dtype = dtype,
+        load_in_4bit = load_in_4bit,
+    )
+    
+    FastLanguageModel.for_inference(model)
+    
+    # evaluate
+    acc = 0
+    for idx, dp in enumerate(test_dataset): # TODO: batch
+        message = dp['conversations'][0]
+        label = dp['conversations'][1]
+        assert label['role'] == "assistant"
+        
+        # format the input
+        inputs = tokenizer.apply_chat_template([message], tokenize = True, add_generation_prompt = True, return_tensors = 'pt').to('cuda')
+        print(f"222, {inputs}")
+        
+        text_streamer = TextStreamer(tokenizer, skip_prompt = True)
+        outputs = model.generate(input_ids = inputs, streamer = text_streamer, max_new_tokens = 128,
+                        use_cache = True, temperature = 1.5, min_p = 0.1)
+        outputs = tokenizer.batch_decode(outputs)
+        parts = outputs[0].split("<|start_header_id|>assistant<|end_header_id|>\n\n")
+        # print(f"111, {parts}")
+        results = parts[1].strip("<|eot_id|>")
+        results = contractions.fix(results)
+        try:
+            results = ast.literal_eval(results)
+        except:
+            try:
+                results = re.sub(r"\['(.*?)'", remove_inner_single_quotes, results) 
+                results = ast.literal_eval(results)
+            except:
+                results = {
+                    "Metapath":"",
+                    "Restriction":{},
+                }
+
+        
+        
+        pred_metapath = results['Metapath']
+        
+        ground_metapath = ast.literal_eval(label['content'])['Metapath']
+        if pred_metapath == ground_metapath:
+            acc += 1
+        print(f"Prediction: {pred_metapath}")
+        print(f"Ground truth: {ground_metapath}")
+        
+        
+    
+    print(f"Accuracy: {acc / len(test_dataset)}")
+    
+
+def main(dataset_name, model_name):
+    # *****load dataset*****
+    data_dir = f"./data/finetune"
+    data_path = os.path.join(data_dir, f"{dataset_name}/llama_ft.jsonl")
+    # data_path = os.path.join(data_dir, f"{dataset_name}/llama_ft_{model_name}.jsonl")
+    dataset = load_dataset("json", data_files=data_path)
+    dataset = dataset['train']
+
+    # Add an index column to keep track of original indices
+    dataset = dataset.add_column("index", list(range(len(dataset))))
+
+    # Perform train-test split
+    train_test = dataset.train_test_split(test_size=0.2, seed=42)
+    val_test = train_test['test'].train_test_split(test_size=0.5, seed=42)
+
+    # Extract the datasets
+    train_dataset = train_test['train']
+    val_dataset = val_test['train']
+    test_dataset = val_test['test']
+
+    # Remove the index column if not needed
+    train_dataset = train_dataset.remove_columns(['index'])
+    val_dataset = val_dataset.remove_columns(['index'])
+    test_dataset = test_dataset.remove_columns(['index'])
+    
+    # *****train and evaluate*****
+    checkpoint_path = train(train_dataset, val_dataset)
+    
+    
+    checkpoint_path = f"./checkpoints/{dataset_name}/lora_model"
+    # checkpoint_path = f"./checkpoints/{dataset_name}/lora_model_{model_name}"
+    evaluate(test_dataset, checkpoint_path)
+
+if __name__ == "__main__":
+    dataset_name_list = ['mag', 'amazon', 'prime']
+    model_names = ["4o"] 
+    for dataset_name in dataset_name_list:
+        for model_name in model_names:        
+            main(dataset_name, model_name)
+            
+            
+            
+            

Planning/model.py

@@ -0,0 +1,89 @@
+import sys
+import os
+import torch.nn as nn
+import ast
+from unsloth import FastLanguageModel
+from transformers import TextStreamer
+import contractions
+import re
+from Planning.utils import remove_inner_single_quotes
+
+
+class Planner(nn.Module):
+    def __init__(self, dataset_name):
+        super(Planner, self).__init__()
+        self.dataset_name = dataset_name
+        self.checkpoint_path = f"Planning/checkpoints/{dataset_name}/lora_model/"
+        self.max_seq_length = 2048
+        self.dtype = None
+        self.load_in_4bit = True
+        
+        model, tokenizer = FastLanguageModel.from_pretrained(
+        model_name = self.checkpoint_path,
+        max_seq_length = self.max_seq_length,
+        dtype = self.dtype,
+        load_in_4bit = self.load_in_4bit
+        )
+
+        FastLanguageModel.for_inference(model)
+
+        self.model = model
+        self.tokenizer = tokenizer
+        
+    def forward(self, query):
+        message = {'content': query, 'role': 'user'}
+        inputs = self.tokenizer.apply_chat_template(
+        [message],
+        tokenize = True,
+        add_generation_prompt = True, # Must add for generation
+        return_tensors = "pt",
+        ).to("cuda")
+
+        text_streamer = TextStreamer(self.tokenizer, skip_prompt = True)
+        outputs = self.model.generate(input_ids = inputs, streamer = text_streamer, max_new_tokens = 128, # max_new_tokens is the maximum number of new tokens generated beyond the input
+                        use_cache = True, temperature = 1.5, min_p = 0.1) # min_p is a cumulative probability, which makes the generation more diverse
+        
+        
+        outputs = self.tokenizer.batch_decode(outputs)
+        parts = outputs[0].split("<|start_header_id|>assistant<|end_header_id|>\n\n")
+        
+        if len(parts) > 1:
+            results = parts[1].replace("<|eot_id|>", "")
+        else:
+            raise ValueError
+
+
+        # ******* special processing for prime dataset
+        if self.dataset_name == 'prime':
+            try:
+                    # Parse the string using ast.literal_eval
+                    parsed_dict = ast.literal_eval(results)
+                    
+                    return parsed_dict
+            except (SyntaxError, ValueError) as e:
+                    print(f"Error parsing the string: {e}")
+                    return {
+                            "Metapath": "",
+                            "Restriction": {}
+                        }
+        
+        
+        results = contractions.fix(results)
+
+        try:
+            results = ast.literal_eval(results)
+        except:
+            print(f"Fail")
+            try:
+                results = re.sub(r"\['(.*?)'", remove_inner_single_quotes, results) # TODO: need optimize
+                results = ast.literal_eval(results)
+            except:
+                results = {
+                "Metapath": "",
+                "Restriction": {},
+                
+            }
+        rg = results
+
+            
+        return rg

Planning/utils.py

@@ -0,0 +1,4 @@
+def remove_inner_single_quotes(match):
+    content = match.group(1)  # Get the string inside ['']
+    cleaned_content = content.replace("'", "")  # Remove inner single quotes
+    return f"['{cleaned_content}"

Reasoning/mor4path.py

@@ -0,0 +1,435 @@
+import sys
+from pathlib import Path
+# Get the absolute path of the current script
+current_file = Path(__file__).resolve()
+project_root = current_file.parents[2]
+# Add the project root to the system path
+sys.path.append(str(project_root))
+
+from .utils import combine_dicts, parse_metapath, get_scorer, get_text_retriever, fix_length
+from models.model import ModelForSTaRKQA
+
+
+class MOR4Path(ModelForSTaRKQA):
+    def __init__(self, dataset_name, text_retriever_name, scorer_name, skb, topk=100):
+        super(MOR4Path, self).__init__(skb)
+        self.dataset_name = dataset_name
+        self.text_retriever = get_text_retriever(dataset_name, text_retriever_name, skb)
+        self.scorer = get_scorer(dataset_name, scorer_name=scorer_name, skb=skb)
+        # self.scorer = self.text_retriever
+        self.topk = topk
+        self.node_type_list = skb.node_type_lst()
+        self.edge_type_list = skb.rel_type_lst()
+        if self.dataset_name == "prime":
+            self.tp_list = skb.get_tuples()
+            self.target_type_list = skb.candidate_types
+        else:
+            self.tp_dict = {(tp[0], tp[-1]): tp[1] for tp in skb.get_tuples()}
+            self.target_type_list = ['paper' if dataset_name == 'mag' else 'product']
+            
+        self.skb = skb
+        self.ini_k = 5 # topk for initial retrieval
+        self.mor_k = 10 # topk for textual retrieval in MOR
+        self.mor_count = 0
+        self.num_negs = 200
+        
+        
+        
+    def rg2routes(self, rg):
+        """
+            input: rg: {"Metapath": "", "Restriction": {}}
+            output: routes: [['paper', 'author', 'paper'], ['paper', 'paper']]
+        """
+        # parse rg
+        metapath = rg["Metapath"]
+        if isinstance(rg["Metapath"], list):
+            routes = rg["Metapath"]
+        elif isinstance(rg["Metapath"], str):
+            routes = parse_metapath(metapath)
+        else:
+            return None
+        
+        return routes
+    
+    def check_valid(self, routes, rg):
+        # check the length of routes
+        if not routes:
+            # raise ValueError(f"Empty routes: {routes}")
+            return None
+        
+        if len(routes) == 1 and len(routes[0]) == 1: # single node, directly do text retrieval
+            return None
+
+        
+        # Step 1: Filter routes by target type
+        target_type_valid_routes = [
+            route for route in routes if route[-1] in self.target_type_list
+        ]
+        if not target_type_valid_routes:
+    
+            return None
+
+        # Step 2: Filter routes by node and edge type
+        type_valid_routes = [
+            route
+            for route in target_type_valid_routes
+            if all(
+                node in self.node_type_list or node in self.edge_type_list
+                for node in route
+            )
+        ]
+        if not type_valid_routes:
+            
+            return None
+
+        # Step 3: Check existence of relations
+        relation_valid_routes = []
+        for route in type_valid_routes:
+            if self.dataset_name == "prime":
+                if len(route) < 3:
+                    continue
+                triplets = [
+                    (route[i], route[i + 1], route[i + 2])
+                    for i in range(0, len(route) - 2, 2)
+                ]
+                
+                if all(tp in self.tp_list for tp in triplets) and all(len(tp) == 3 for tp in triplets): # and all length of triplets is 3
+                    relation_valid_routes.append(route)
+            else:
+                pairs = [(route[i], route[i + 1]) for i in range(len(route) - 1)]
+                if all(tp in self.tp_dict.keys() for tp in pairs):
+                    relations = [self.tp_dict[tp] for tp in pairs]
+                    
+                    # make route with relations
+                    new_route = []
+                    for i in range(len(relations)):
+                        new_route.append(pairs[i][0])
+                        new_route.append(relations[i])
+                    new_route.append(pairs[-1][-1])
+                    
+                    relation_valid_routes.append(new_route)             
+
+        if not relation_valid_routes:
+            
+            return None
+
+        return relation_valid_routes
+    
+    def get_candidates4route(self, query, q_id, route, restriction):
+        # initialization
+        
+        ini_node_type = route[0]
+        
+        try:
+            type_restr = "".join(restriction[ini_node_type])
+        except:
+            type_restr = ""
+            
+        ini_dict = self.text_retriever.retrieve(query + " " + type_restr, q_id=q_id, topk=self.ini_k, node_type=ini_node_type)
+        current_node_ids = list(ini_dict.keys())
+        
+        
+        # initialize the bm_vector_dict
+        bm_vector_dict = {key: [value] for key, value in ini_dict.items()}
+        
+        # initilization for paths
+        paths = {}
+        for c_id in current_node_ids:
+            paths[c_id] = [c_id]
+    
+        # loop
+        hops = len(route)
+        # for hop/layer
+        for hop in range(0, hops-2, 2):
+            new_paths = {}
+            
+            cur_node_type = route[hop]
+            next_node_type = route[hop+2]
+            edge_type = route[hop+1]
+            next_node_ids = []
+            
+            new_vector_dict = {}
+            # for node
+            for node_id in current_node_ids:
+                neighbor_ids = self.skb.get_neighbor_nodes(idx=node_id, edge_type=edge_type)
+                next_node_ids.extend(neighbor_ids)
+                
+                # ***** update paths and score_vector_dict *****
+                for neighbor_id in neighbor_ids:
+                    if neighbor_id not in new_paths.keys(): # only add new node
+                        new_paths[neighbor_id] = paths[node_id] + [neighbor_id]
+                        new_vector_dict[neighbor_id] = bm_vector_dict[node_id] + [-1] # -1 for padding
+                    
+    
+            bm_vector_dict = new_vector_dict 
+            
+            
+            
+            # ***** layer text retrieval *****
+            # if there is restriction for the next node, add text_retriever
+            if next_node_type in restriction.keys() and len(restriction[next_node_type]) > 0 and restriction[next_node_type] != [""]:
+                try:
+
+                    retrieve_dict = self.text_retriever.retrieve(query+" "+"".join(restriction[next_node_type]), q_id=q_id, topk=self.mor_k, node_type=route[hop+2])
+
+                    new_query = query+ " " + "".join(restriction[next_node_type])
+                    
+                    # take union
+                    next_node_ids.extend(list(set(retrieve_dict.keys())))
+                    
+                    # ***** update paths and bm_vector_dict *****
+                    for c_id in retrieve_dict.keys():
+                        if c_id not in new_paths.keys():
+                            new_paths[c_id] = [c_id]
+                            bm_vector_dict[c_id] = [retrieve_dict[c_id]]
+                            
+                except:
+                    pass
+        
+                
+            paths = new_paths
+            current_node_ids = list(set(next_node_ids))
+
+
+        candidates = current_node_ids
+        
+        self.paths.append(paths)
+        self.bm_vector_dict.append(bm_vector_dict)
+        
+        
+        return candidates
+    
+    def merge_candidate_pools(self, non_empty_candidates_lists):
+         
+         # if only one non-empy candidates list left, return it as a set
+         if len(non_empty_candidates_lists) == 1:
+              return set(non_empty_candidates_lists[0])   
+          
+         # find the intersection candidates ids
+         result = set(non_empty_candidates_lists[0])
+         for lst in non_empty_candidates_lists[1:]:
+              result.intersection_update(lst)
+        
+         # if the intersection is empty, return the union of all candidates
+         if len(result) == 0:
+              result = set()
+              for lst in non_empty_candidates_lists:
+                   result.update(lst)        
+         
+         return list(result)
+    
+    def get_mor_candidates(self, query, q_id, valid_routes, restriction):
+        
+        # Step 1: Get candidates for each route
+        candidates_pool = []
+        for route in valid_routes:
+            if route[0] in restriction.keys() and len(restriction[route[0]]) > 0:
+                candidates_pool.append(self.get_candidates4route(query, q_id, route, restriction)) # topk is the candidates retrieved from textual retriever
+        
+        
+        # remove empty lists from candidates
+        non_empty_candidates_lists = [lst for lst in candidates_pool if lst]
+        if len(non_empty_candidates_lists) == 0:
+
+            return {}
+        
+        
+        # Step 2: Combine candidates from different routes, try intersection first, then union
+        candidates = self.merge_candidate_pools(non_empty_candidates_lists) # candidates is a list
+        
+        # step 3: score the candidates, ini to -1
+        pred_dict = dict(zip(candidates, [-1]*len(candidates)))
+        # print(f"111, {pred_dict}")
+        
+        return pred_dict
+    
+    def check_topk(self, query, q_id, pred_dict):
+        
+        missing = self.topk - len(set(pred_dict.keys()))
+        if missing > 0:
+            added_dict = self.text_retriever.retrieve(query, q_id, topk=self.topk+20, node_type=self.target_type_list) # +20 make it more safe 
+            available_nodes = {key: value for key, value in added_dict.items() if key not in pred_dict.keys()}
+            sorted_available_nodes = sorted(available_nodes.items(), key=lambda x: x[1], reverse=True)
+            # Select only the required number of nodes to fill the missing slots
+            selected_nodes = dict(sorted_available_nodes[:missing])
+            
+            # Update pred_dict with the selected nodes
+            pred_dict.update(selected_nodes)
+            
+            # updata paths
+            for node_id in selected_nodes.keys():
+                self.paths[node_id] = [node_id]
+                
+            # update bm_vector_dict
+            new_bm_vector_dict = {key: [value] for key, value in selected_nodes.items()}
+            self.bm_vector_dict.update(new_bm_vector_dict)
+            
+        scored_dict = self.scorer.score(query, q_id=q_id, candidate_ids=list(pred_dict.keys()))
+        
+        if len(scored_dict) > self.topk:
+            # initiliaze the new_paths
+            new_paths = {}
+            
+            # Select the top-k nodes based on the scores
+            sorted_scored_dict = sorted(scored_dict.items(), key=lambda x: x[1], reverse=True)
+            scored_dict = dict(sorted_scored_dict[:self.topk])
+            
+            # update paths
+            for node_id in scored_dict.keys():
+                new_paths[node_id] = self.paths[node_id]
+            
+            self.paths = new_paths
+            
+            # update bm_vector_dict
+            new_bm_vector_dict = {node_id: self.bm_vector_dict[node_id] for node_id in scored_dict.keys()}
+            self.bm_vector_dict = new_bm_vector_dict
+            
+        
+        return scored_dict
+    
+    
+    # check fixed negtopk
+    def check_negtopk(self, query, q_id, pred_dict, ans_ids):
+        # check the positive nodes
+        pos_ids = [node_id for node_id in ans_ids if node_id in pred_dict.keys()]
+        pos_dict = {key: value for key, value in pred_dict.items() if key in pos_ids}
+        neg_ids = pred_dict.keys() - set(pos_ids)
+        neg_dict = {key: value for key, value in pred_dict.items() if key in neg_ids}
+        
+        # check the number of negative nodes
+        missing = self.num_negs - len(neg_ids)
+        
+        if missing > 0:
+            
+            added_dict = self.text_retriever.retrieve(query, q_id, topk=self.num_negs+200, node_type=self.target_type_list) # +20 make it more safe 
+            available_nodes = {key: value for key, value in added_dict.items() if key not in pred_dict.keys() and key not in ans_ids}
+            sorted_available_nodes = sorted(available_nodes.items(), key=lambda x: x[1], reverse=True)
+            # Select only the required number of nodes to fill the missing slots
+            selected_nodes = dict(sorted_available_nodes[:missing])
+            
+            # Update pred_dict with the selected nodes
+            neg_dict.update(selected_nodes)
+            
+            # updata paths
+            for node_id in selected_nodes.keys():
+                self.paths[node_id] = [node_id]
+                
+            # update bm_vector_dict
+            new_bm_vector_dict = {key: [value] for key, value in selected_nodes.items()}
+            self.bm_vector_dict.update(new_bm_vector_dict)
+            
+            
+        scored_neg_dict = self.scorer.score(query, q_id=q_id, candidate_ids=list(neg_dict.keys()))
+        if pos_dict:
+            scored_pos_dict = self.scorer.score(query, q_id=q_id, candidate_ids=list(pos_dict.keys()))
+        else:
+            scored_pos_dict = {}
+        
+        if len(scored_neg_dict) > self.num_negs:
+            # Select the top-k nodes based on the scores
+            sorted_scored_neg_dict = sorted(scored_neg_dict.items(), key=lambda x: x[1], reverse=True)
+            scored_neg_dict = dict(sorted_scored_neg_dict[:self.num_negs])
+            
+            
+            
+        scored_neg_dict.update(scored_pos_dict)
+        scored_dict = scored_neg_dict
+        print(len(scored_dict))
+        
+        # update paths
+        new_paths = {}
+        for node_id in scored_dict.keys():
+            new_paths[node_id] = self.paths[node_id]
+        self.paths = new_paths
+        
+        # update bm_vector_dict
+        new_bm_vector_dict = {node_id: self.bm_vector_dict[node_id] for node_id in scored_dict.keys()}
+        self.bm_vector_dict = new_bm_vector_dict
+        
+        
+        return scored_dict
+        
+    
+    def forward(self, query, q_id, ans_ids, rg, args):
+        
+        self.paths = []
+        self.bm_vector_dict = []
+        self.ada_score = {}
+        # ***** Structural Retrieval *****
+        
+        # reasoning grpah to routes
+        if self.dataset_name == "prime":
+            routes = rg["Metapath"]
+        else:
+            routes = self.rg2routes(rg)
+        
+        # check valid
+        valid_routes = self.check_valid(routes, rg) # add check for restriction
+        
+        if valid_routes is None:
+            # do textual retrieval
+            pred_dict = self.text_retriever.retrieve(query, q_id, topk=self.topk, node_type=self.target_type_list)  
+            
+            # update bm_vector_dict
+            self.bm_vector_dict = {key: [value] for key, value in pred_dict.items()}
+            
+        else:
+            # truncate the valid_routes
+            if self.dataset_name == "prime":
+                pass
+            else:
+                valid_routes = [route[-5:] for route in valid_routes]
+                
+            
+            # do structural retrieval
+            restriction = rg["Restriction"]
+            pred_dict = self.get_mor_candidates(query, q_id, valid_routes, restriction)
+            self.mor_count += 1
+            
+        
+        # **** combine paths ****
+        if self.paths:
+            self.paths = combine_dicts(self.paths, pred_dict=pred_dict) # return dict
+        else:
+            self.paths = {}
+            for node_id in pred_dict.keys():
+                self.paths[node_id] = [node_id]
+                
+        # ***** combine bm_vector_dict *****
+        if isinstance(self.bm_vector_dict, list):
+            self.bm_vector_dict = combine_dicts(self.bm_vector_dict, pred_dict=pred_dict)
+        
+        # **** fix neg for training; fix candidates for testing ****
+        if args.mod == "train":
+            # check neg topk 
+            pred_dict = self.check_negtopk(query, q_id, pred_dict, ans_ids)
+        else:
+            # check topk    
+            pred_dict = self.check_topk(query, q_id, pred_dict)
+
+        
+        # **** length padding and truncate *****
+        if self.dataset_name != "prime":
+            self.paths = fix_length(self.paths)
+                
+        if len(self.paths) != len(pred_dict):
+            print(f"paths: {self.paths}")
+            print(f"pred_dict: {pred_dict}")
+            raise ValueError(f"Length mismatch between paths and pred_dict: {len(self.paths)}, {len(pred_dict)}")
+
+        output = {
+            "query": query,
+            "pred_dict": pred_dict,
+            "ans_ids": ans_ids,
+            'paths': self.paths,
+            'bm_vector_dict': self.bm_vector_dict,
+            'rg': rg
+        }
+        
+        
+        return output
+    
+if __name__ == "__main__":
+    print(f"Test mor4path")
+        
+

Reasoning/structural_retriever/stru4path.py

@@ -0,0 +1,305 @@
+"""
+input:   rg
+output (fixed 100 candidates, for path-based reranking): 
+{
+    "query": query,
+    "pred_dict": {node_id: score},
+    "ans_ids": [],
+    'paths': {node_id: [node_ids_path]}
+}
+
+"""
+import sys
+import os
+sys.path.append(os.path.dirname(os.path.abspath(os.getcwd())))
+
+from utils import combine_dicts, parse_metapath, get_scorer, get_text_retriever, fix_length
+from models.model import ModelForSTaRKQA
+import time
+
+
+
+class Stru4Path(ModelForSTaRKQA):
+    def __init__(self, dataset_name, text_retriever_name, scorer_name, skb, topk=100):
+        super(Stru4Path, self).__init__(skb)
+        self.dataset_name = dataset_name
+        self.text_retriever = get_text_retriever(dataset_name, text_retriever_name, skb)
+        self.scorer = get_scorer(dataset_name, scorer_name=scorer_name, skb=skb)
+        # self.scorer = self.text_retriever
+        self.topk = topk
+        self.node_type_list = skb.node_type_lst()
+        self.edge_type_list = skb.rel_type_lst()
+        if self.dataset_name == "prime":
+            self.tp_list = skb.get_tuples()
+            self.target_type_list = skb.candidate_types
+        else:
+            self.tp_dict = {(tp[0], tp[-1]): tp[1] for tp in skb.get_tuples()}
+            self.target_type_list = ['paper' if dataset_name == 'mag' else 'product']
+            
+        self.skb = skb
+        self.ini_k = 5 # topk for initial retrieval
+        self.stru_count = 0
+    
+        
+        
+        
+    def rg2routes(self, rg):
+        """
+            input: rg: {"Metapath": "", "Restriction": {}}
+            output: routes: [['paper', 'author', 'paper'], ['paper', 'paper']]
+        """
+        # parse rg
+        metapath = rg["Metapath"]
+        if isinstance(rg["Metapath"], list):
+            routes = rg["Metapath"]
+        elif isinstance(rg["Metapath"], str):
+            routes = parse_metapath(metapath)
+        else:
+            return None
+        
+        return routes
+    
+    def check_valid(self, routes, rg):
+        # check the length of routes
+        if not routes:
+            # raise ValueError(f"Empty routes: {routes}")
+            return None
+        
+        if len(routes) == 1 and len(routes[0]) == 1: # single node, directly do text retrieval
+            return 1
+        
+        # Step 1: Filter routes by target type
+        target_type_valid_routes = [
+            route for route in routes if route[-1] in self.target_type_list
+        ]
+        if not target_type_valid_routes:
+            return None
+
+        # Step 2: Filter routes by node and edge type
+        type_valid_routes = [
+            route
+            for route in target_type_valid_routes
+            if all(
+                node in self.node_type_list or node in self.edge_type_list
+                for node in route
+            )
+        ]
+        if not type_valid_routes:
+            return None
+
+        # Step 3: Check existence of relations
+        relation_valid_routes = []
+        for route in type_valid_routes:
+            if self.dataset_name == "prime":
+                triplets = [
+                    (route[i], route[i + 1], route[i + 2])
+                    for i in range(0, len(route) - 2, 2)
+                ]
+                
+                if all(tp in self.tp_list for tp in triplets):
+                    relation_valid_routes.append(route)
+            else:
+                pairs = [(route[i], route[i + 1]) for i in range(len(route) - 1)]
+                if all(tp in self.tp_dict.keys() for tp in pairs):
+                    relations = [self.tp_dict[tp] for tp in pairs]
+                    
+                    # make route with relations
+                    new_route = []
+                    for i in range(len(relations)):
+                        new_route.append(pairs[i][0])
+                        new_route.append(relations[i])
+                    new_route.append(pairs[-1][-1])
+                    # print(f"222, {new_route}")
+                    
+                    relation_valid_routes.append(new_route)             
+
+        if not relation_valid_routes:
+            return None
+
+        return relation_valid_routes
+    
+    def get_candidates4route(self, query, q_id, route, restriction):
+        # initialization
+        
+        ini_node_type = route[0]
+        
+        try:
+            extra_restr = "".join(restriction[ini_node_type])
+        except:
+            extra_restr = ""
+        ini_dict = self.text_retriever.retrieve(query + " " + extra_restr, q_id=q_id, topk=self.ini_k, node_type=ini_node_type)
+        current_node_ids = list(ini_dict.keys())
+        
+        # initilization for paths
+        paths = {}
+        for c_id in current_node_ids:
+            paths[c_id] = [c_id]
+    
+        # loop
+        hops = len(route)
+        # for hop/layer
+        for hop in range(0, hops-2, 2):
+            new_paths = {}
+            
+            cur_node_type = route[hop]
+            next_node_type = route[hop+2]
+            edge_type = route[hop+1]
+            next_node_ids = []
+            
+            # for node
+            for node_id in current_node_ids:
+                neighbor_ids = self.skb.get_neighbor_nodes(idx=node_id, edge_type=edge_type)
+                next_node_ids.extend(neighbor_ids)
+                
+                # **x*** update paths *****
+                for neighbor_id in neighbor_ids:
+                    new_paths[neighbor_id] = paths[node_id] + [neighbor_id]
+        
+                        
+            paths = new_paths
+            
+            current_node_ids = list(set(next_node_ids))
+
+        candidates = current_node_ids
+        self.paths.append(paths)
+        
+        
+        return candidates
+    
+    def merge_candidate_pools(self, non_empty_candidates_lists):
+         
+         
+         # if only one non-empy candidates list left, return it as a set
+         if len(non_empty_candidates_lists) == 1:
+              return set(non_empty_candidates_lists[0])   
+         # find the intersection candidates ids
+         result = set(non_empty_candidates_lists[0])
+         for lst in non_empty_candidates_lists[1:]:
+              result.intersection_update(lst)
+        
+         # if the intersection is empty, return the union of all candidates
+         if len(result) == 0:
+              result = set()
+              for lst in non_empty_candidates_lists:
+                   result.update(lst)        
+         
+         
+         
+         return list(result)
+    
+    def get_mor_candidates(self, query, q_id, valid_routes, restriction):
+        
+        # Step 1: Get candidates for each route
+        candidates_pool = []
+        for route in valid_routes:
+            if route[0] in restriction.keys() and len(restriction[route[0]]) > 0:
+                candidates_pool.append(self.get_candidates4route(query, q_id, route, restriction)) # topk is the candidates retrieved from textual retriever    
+        
+        non_empty_candidates_lists = [lst for lst in candidates_pool if lst]
+        if not non_empty_candidates_lists: # no candidates, return empty dict
+            print(f"123, {non_empty_candidates_lists}")
+            
+            # raise ValueError("No candidates for any route")
+            return {}
+        
+        
+        # Step 2: Combine candidates from different routes, try intersection first, then union
+        candidates = self.merge_candidate_pools(candidates_pool) # candidates is a list
+        if not candidates:
+            return {}
+        
+        
+        # step 3: score the candidates, ini to -1
+        pred_dict = dict(zip(candidates, [-1]*len(candidates)))
+        # print(f"111, {pred_dict}")
+        
+        return pred_dict
+    
+        
+    
+    def forward(self, query, q_id, ans_ids, rg):
+        
+        self.paths = []
+        # ***** Structural Retrieval *****
+        
+        # reasoning grpah to routes
+        s_time = time.time()
+        routes = self.rg2routes(rg)
+        # print(f"444, {time.time()-s_time}")
+        
+        # check valid
+        s_time = time.time()
+        valid_routes = self.check_valid(routes, rg) # add check for restriction
+        # print(f"555, {time.time()-s_time}")
+        
+        if valid_routes is None:
+            # return empty dict
+            return {
+                "query": query,
+                "pred_dict": {},
+                "ans_ids": ans_ids,
+                'paths': {},
+                'query_pattern': rg['Metapath']
+            }
+        elif valid_routes == 1: # TODO: empty string
+            print(f"1234: {valid_routes}")
+            # do text retrieval
+            pred_dict = self.text_retriever.retrieve(query, q_id=q_id, topk=self.topk, node_type=f'{self.target_type_list[0]}')
+            
+        else:
+            # do structural retrieval
+            # truncate the valid_routes
+            if self.dataset_name == "prime":
+                pass
+            else:
+                valid_routes = [route[-5:] for route in valid_routes]
+            
+            restriction = rg["Restriction"]
+            pred_dict = self.get_mor_candidates(query, q_id, valid_routes, restriction)
+            self.stru_count += 1
+        
+        # **** combine paths ****
+        if self.paths:
+            self.paths = combine_dicts(self.paths, pred_dict=pred_dict) # return dict
+            
+        else:
+            self.paths = {}
+            for node_id in pred_dict.keys():
+                self.paths[node_id] = [node_id]
+        
+        # if retrieved candidates is empty, return empty dict
+        if not pred_dict:
+            return {
+                "query": query,
+                "pred_dict": {},
+                "ans_ids": ans_ids,
+                'paths': {},
+                'query_pattern': rg['Metapath']
+            }
+        
+        # score the candidates
+        pred_dict = self.scorer.score(query, q_id, list(pred_dict.keys()))
+        
+        # # **** length padding and truncate *****
+        # self.paths = fix_length(self.paths)
+                
+        if len(self.paths) != len(pred_dict):
+            print(f"paths: {self.paths}")
+            print(f"pred_dict: {pred_dict}")
+            raise ValueError(f"Length mismatch between paths and pred_dict: {len(self.paths)}, {len(pred_dict)}")
+
+        output = {
+            "query": query,
+            "pred_dict": pred_dict,
+            "ans_ids": ans_ids,
+            'paths': self.paths,
+            'query_pattern': rg['Metapath'],
+            'rg': rg
+        }
+        
+        
+        return output
+        
+
+        
+

Reasoning/text_retrievers/__init__.py

@@ -0,0 +1,4 @@
+from .stark_model import ModelForSTaRKQA
+from .ada import Ada
+from .bm25 import BM25
+from .contriever import Contriever

Reasoning/text_retrievers/ada.py

@@ -0,0 +1,66 @@
+
+"""
+input: query, query_id, candidates_ids
+output: pred_dict: {node_id: similarity}
+"""
+
+import sys
+from pathlib import Path
+# Get the absolute path of the current script
+current_file = Path(__file__).resolve()
+project_root = current_file.parents[2]
+# Add the project root to the system path
+sys.path.append(str(project_root))
+import torch
+from Reasoning.text_retrievers.stark_model import ModelForSTaRKQA
+
+class Ada(ModelForSTaRKQA):
+    def __init__(self, skb, dataset_name, device):
+        super(Ada, self).__init__(skb)
+        self.emb_dir = f"{project_root}/Reasoning/data/emb/{dataset_name}/"
+        self.query_emb_path = self.emb_dir + "text-embedding-ada-002/query/query_emb_dict.pt"
+        self.query_emb_dict = torch.load(self.query_emb_path)
+        # print(f"777, {self.query_emb_path}")
+
+        self.candidate_emb_path = self.emb_dir + "text-embedding-ada-002/doc/candidate_emb_dict.pt"
+        self.candidate_emb_dict = torch.load(self.candidate_emb_path)
+        self.device = device
+
+        assert len(self.candidate_emb_dict) == len(self.candidate_ids)
+
+        candidate_embs = [self.candidate_emb_dict[idx].view(1, -1) for idx in self.candidate_ids]
+        self.candidate_embs = torch.cat(candidate_embs, dim=0).to(device)
+
+    def score(self, query, q_id, candidate_ids):
+        """
+            pred_dict[node_id] = similarity (tensor)
+
+        """
+        # Dimension of query_emb: torch.Size([1, emb_dim])
+        query_emb = self.query_emb_dict[q_id].view(1, -1)
+        # Dimension of candidates_embs: torch.Size([num_candidates, emb_dim])
+        # # candidates_embs = self.candidate_embs[candidates_ids]
+        candi_embs = [self.candidate_emb_dict[c_id].view(1, -1) for c_id in candidate_ids]
+        candidates_embs = torch.cat(candi_embs, dim=0).to(self.device)
+        # Dimension of similarity: torch.Size([num_candidates])
+        similarity = torch.matmul(query_emb.to(self.device), candidates_embs.T).squeeze(dim=0).cpu()
+        pred_dict = {}
+        for i in range(len(candidate_ids)):
+            pred_dict[candidate_ids[i]] = similarity[i].item()
+        
+        return pred_dict
+    
+    def retrieve(self, query, q_id, topk, node_type=None):
+        # Dimension of query_emb: torch.Size([1, emb_dim])
+        query_emb = self.query_emb_dict[q_id].view(1, -1)
+
+        similarity = torch.matmul(query_emb.to(self.device), self.candidate_embs.T).cpu()
+        if isinstance(query, str):
+            pred_dict = dict(zip(self.candidate_ids, similarity.view(-1)))
+
+        sorted_pred_ids = sorted(pred_dict, key=lambda x: pred_dict[x], reverse=True)
+        selected_pred_ids = sorted_pred_ids[:topk]
+        pred_dict = {id: pred_dict[id].item() for id in selected_pred_ids}
+        print(f"sorted: {pred_dict}")
+
+        return pred_dict

Reasoning/text_retrievers/bm25.py

@@ -0,0 +1,108 @@
+"""
+input: query, node_type, topk
+output: pred_dict: {node_id: score}
+
+"""
+import bm25s
+from tqdm import tqdm
+import sys
+from pathlib import Path
+# Get the absolute path of the current script
+current_file = Path(__file__).resolve()
+project_root = current_file.parents[2]
+# Add the project root to the system path
+sys.path.append(str(project_root))
+
+from Reasoning.text_retrievers.stark_model import ModelForSTaRKQA
+
+target_type = {'amazon': 'product', 'prime': 'combine', 'mag': 'paper'}
+
+class BM25(ModelForSTaRKQA):
+    
+    def __init__(self, skb, dataset_name):
+        super(BM25, self).__init__(skb)
+        self.retrievers = {}
+        self.text_to_ids = {}
+        type_names = skb.node_type_lst()
+        self.nodeid_to_index = {}
+        
+        self.target_type = target_type[dataset_name]
+
+        if self.target_type not in type_names:
+             ids = skb.get_candidate_ids()
+             
+             
+             corpus = [skb.get_doc_info(id) for id in tqdm(ids, desc=f"Gathering docs for combine")]     
+             retriever = bm25s.BM25(corpus=corpus)
+             retriever.index(bm25s.tokenize(corpus))
+             # Build hash map from text to node_id
+             text_to_id = {hash(text): id for text, id in zip(corpus, ids)}
+             # Store the retriever and text_to_id by type_name
+             self.retrievers[self.target_type] = retriever
+             self.text_to_ids[self.target_type] = text_to_id
+             
+             self.nodeid_to_index[self.target_type] = {id: i for i, id in enumerate(ids)}
+
+        # Initialize retrievers and text-to-index maps for each type_name
+        for type_name in type_names:
+            ids = skb.get_node_ids_by_type(type_name)
+            
+            # we manually replace '&' with '_and_' to avoid the error in BM25, because BM25 uses '&' as a special character and will not tokenize it
+            corpus = [skb.get_doc_info(id).replace('&', '_and_').replace('P.O.R', 'P_dot_O_dot_R') for id in tqdm(ids, desc=f"Gathering docs for {type_name}")]
+            # Create the BM25 model for the current type_name
+            retriever = bm25s.BM25(corpus=corpus)
+            retriever.index(bm25s.tokenize(corpus))
+            
+            # Build hash map from text to index
+            text_to_id = {hash(text): id for text, id in zip(corpus, ids)}
+
+            # Store the retriever and text_to_id by type_name
+            self.retrievers[type_name] = retriever
+            self.text_to_ids[type_name] = text_to_id
+            
+            # build map from node_id to index
+            self.nodeid_to_index[type_name] = {id: i for i, id in enumerate(ids)}
+    
+    def score(self, query, q_id, candidate_ids):
+        pred_dict = {}
+        
+        for c_id in candidate_ids:
+            type_name = self.skb.get_node_type_by_id(c_id)
+            score = self.retrievers[type_name].get_scores(list(bm25s.tokenize(query)[1].keys()))[self.nodeid_to_index[type_name][c_id]] # save the query tokens 
+            pred_dict[c_id] = score
+            
+        # print(f"999, {pred_dict}")
+        
+        return pred_dict
+        
+    def retrieve(self, query, q_id, topk, node_type=None):
+                
+        """
+        Forward pass to compute similarity scores for the given query.
+
+        Args:
+            query (str): Query string.
+
+        Returns:
+            pred_dict (dict): A dictionary of candidate ids and their corresponding similarity scores.
+        """
+        if '&' in query:
+               query = query.replace('&', '_and_')
+        if 'P.O.R' in query:
+               query = query.replace('P.O.R', 'P_dot_O_dot_R')
+        if isinstance(node_type, list):
+            if len(node_type) > 1:
+                node_type = 'combine'
+            else:
+                node_type = node_type[0]
+        results, scores = self.retrievers[node_type].retrieve(bm25s.tokenize(query), k=topk)
+        ids = [self.text_to_ids[node_type][hash(result.item())] for result in results[0]]
+        scores = scores[0].tolist()
+        pred_dict = dict(zip(ids, scores))
+        # print(f"666, {pred_dict}")
+        
+        return pred_dict
+
+if __name__ == '__main__':
+    print(f"Testing BM25")
+    

Reasoning/text_retrievers/contriever.py

@@ -0,0 +1,78 @@
+
+"""
+input: query, query_id, candidates_ids
+output: pred_dict: {node_id: similarity}
+"""
+import heapq
+import sys
+from pathlib import Path
+# Get the absolute path of the current script
+current_file = Path(__file__).resolve()
+project_root = current_file.parents[2]
+# Add the project root to the system path
+sys.path.append(str(project_root))
+from stark_qa.tools.api_lib.openai_emb import get_contriever, get_contriever_embeddings
+import torch
+from Reasoning.text_retrievers.stark_model import ModelForSTaRKQA
+
+class Contriever(ModelForSTaRKQA):
+    def __init__(self, skb, dataset_name, device):
+        super(Contriever, self).__init__(skb)
+        self.emb_dir = f"{project_root}/Reasoning/data/emb/{dataset_name}/"
+
+        self.query_emb_path = self.emb_dir + "contriever/query_no_rel_no_compact/query_emb_dict.pt"
+        self.query_emb_dict = torch.load(self.query_emb_path)
+
+        self.candidate_emb_path = self.emb_dir + "contriever/doc_no_rel_no_compact/candidate_emb_dict.pt"
+        self.candidate_emb_dict = torch.load(self.candidate_emb_path)
+        self.device = device
+
+        assert len(self.candidate_emb_dict) == len(self.candidate_ids)
+
+        candidate_embs = [self.candidate_emb_dict[idx].view(1, -1) for idx in self.candidate_ids]
+        self.candidate_embs = torch.cat(candidate_embs, dim=0).to(device)
+        
+        # load contriever for query embeddings
+        self.encoder, self.tokenizer = get_contriever(dataset_name=dataset_name)    
+        self.encoder = self.encoder.to(device)    
+        
+
+    def score(self, query, q_id, candidate_ids):
+        """
+            pred_dict[node_id] = similarity (tensor)
+
+        """
+        
+        
+        # Dimension of query_emb: torch.Size([1, emb_dim])
+        query_emb = self.query_emb_dict[q_id].view(1, -1)
+
+        # Dimension of candidates_embs: torch.Size([num_candidates, emb_dim])
+        candi_embs = [self.candidate_emb_dict[c_id].view(1, -1) for c_id in candidate_ids]
+        candidates_embs = torch.cat(candi_embs, dim=0).to(self.device)
+        # Dimension of similarity: torch.Size([num_candidates])
+        similarity = torch.matmul(query_emb.to(self.device), candidates_embs.T).squeeze(dim=0).cpu()
+        pred_dict = {}
+        for i in range(len(candidate_ids)):
+            pred_dict[candidate_ids[i]] = similarity[i].item()
+        
+        return pred_dict
+    
+    def retrieve(self, query, q_id, topk, node_type=None):
+        # Dimension of query_emb: torch.Size([1, emb_dim])
+        query_emb = get_contriever_embeddings(query, encoder=self.encoder, tokenizer=self.tokenizer, device=self.device)
+        
+        similarity = torch.matmul(query_emb.to(self.device), self.candidate_embs.T).cpu()
+        
+        if isinstance(query, str):
+            pred_dict = dict(zip(self.candidate_ids, similarity.view(-1)))
+        
+    
+        selected_pred_ids = heapq.nlargest(topk, pred_dict, key=pred_dict.get)
+        pred_dict = {id: pred_dict[id].item() for id in selected_pred_ids}
+
+
+        return pred_dict
+
+if __name__ == '__main__':
+    print("Testing Contriever")

Reasoning/text_retrievers/stark_model.py

@@ -0,0 +1,151 @@
+import os
+import os.path as osp
+from typing import Any, Union, List, Dict
+import torch
+import torch.nn as nn
+    
+import sys
+from pathlib import Path
+# Get the absolute path of the current script
+current_file = Path(__file__).resolve()
+project_root = current_file.parents[2]
+# Add the project root to the system path
+sys.path.append(str(project_root))
+
+from stark_qa.tools.api import get_api_embeddings, get_sentence_transformer_embeddings
+from stark_qa.tools.local_encoder import get_llm2vec_embeddings, get_gritlm_embeddings
+from stark_qa.evaluator import Evaluator
+
+
+class ModelForSTaRKQA(nn.Module):
+    
+    def __init__(self, skb, query_emb_dir='.'):
+        """
+        Initializes the model with the given knowledge base.
+        
+        Args:
+            skb: Knowledge base containing candidate information.
+        """
+        super(ModelForSTaRKQA, self).__init__()
+        self.skb = skb
+
+        self.candidate_ids = skb.candidate_ids
+        self.num_candidates = skb.num_candidates
+        self.query_emb_dir = query_emb_dir
+
+        query_emb_path = osp.join(self.query_emb_dir, 'query_emb_dict.pt')
+        if os.path.exists(query_emb_path):
+            print(f'Load query embeddings from {query_emb_path}')
+            self.query_emb_dict = torch.load(query_emb_path)
+        else:
+            self.query_emb_dict = {}
+        self.evaluator = Evaluator(self.candidate_ids)
+    
+    def forward(self, 
+                query: Union[str, List[str]], 
+                candidates: List[int] = None,
+                query_id: Union[int, List[int]] = None,
+                **kwargs: Any) -> Dict[str, Any]:
+        """
+        Forward pass to compute predictions for the given query.
+        
+        Args:
+            query (Union[str, list]): Query string or a list of query strings.
+            candidates (Union[list, None]): A list of candidate ids (optional).
+            query_id (Union[int, list, None]): Query index (optional).
+            
+        Returns:
+            pred_dict (dict): A dictionary of predicted scores or answer ids.
+        """
+        raise NotImplementedError
+    
+    def get_query_emb(self, 
+                      query: Union[str, List[str]], 
+                      query_id: Union[int, List[int]], 
+                      emb_model: str = 'text-embedding-ada-002', 
+                      **encode_kwargs) -> torch.Tensor:
+        """
+        Retrieves or computes the embedding for the given query.
+        
+        Args:
+            query (str): Query string.
+            query_id (int): Query index.
+            emb_model (str): Embedding model to use.
+            
+        Returns:
+            query_emb (torch.Tensor): Query embedding.
+        """        
+        if isinstance(query_id, int):
+            query_id = [query_id]
+        if isinstance(query, str):
+            query = [query]
+
+        encode_kwargs['is_query'] = True
+        if query_id is None:
+            query_emb = get_embeddings(query, emb_model, **encode_kwargs)
+        elif set(query_id).issubset(set(list(self.query_emb_dict.keys()))):
+            query_emb = torch.concat([self.query_emb_dict[qid] for qid in query_id], dim=0)
+        else:
+            query_emb = get_embeddings(query, emb_model, **encode_kwargs)
+            for qid, emb in zip(query_id, query_emb):
+                self.query_emb_dict[qid] = emb.view(1, -1)
+            torch.save(self.query_emb_dict, osp.join(self.query_emb_dir, 'query_emb_dict.pt'))
+            
+        query_emb = query_emb.view(len(query), -1)
+        return query_emb
+    
+    def evaluate(self, 
+                 pred_dict: Dict[int, float], 
+                 answer_ids: Union[torch.LongTensor, List[Any]], 
+                 metrics: List[str] = ['mrr', 'hit@3', 'recall@20'], 
+                 **kwargs: Any) -> Dict[str, float]:
+        """
+        Evaluates the predictions using the specified metrics.
+        
+        Args:
+            pred_dict (Dict[int, float]): Predicted answer ids or scores.
+            answer_ids (torch.LongTensor): Ground truth answer ids.
+            metrics (List[str]): A list of metrics to be evaluated, including 'mrr', 'hit@k', 'recall@k', 
+                                 'precision@k', 'map@k', 'ndcg@k'.
+                             
+        Returns:
+            Dict[str, float]: A dictionary of evaluation metrics.
+        """
+        return self.evaluator(pred_dict, answer_ids, metrics)
+    
+    def evaluate_batch(self, 
+                pred_ids: List[int],
+                pred: torch.Tensor, 
+                answer_ids: Union[torch.LongTensor, List[Any]], 
+                metrics: List[str] = ['mrr', 'hit@3', 'recall@20'], 
+                **kwargs: Any) -> Dict[str, float]:
+        return self.evaluator.evaluate_batch(pred_ids, pred, answer_ids, metrics)
+
+
+def get_embeddings(text, model_name, **encode_kwargs):
+    """
+    Get embeddings for the given text using the specified model.
+    
+    Args:
+        model_name (str): Model name.
+        text (Union[str, List[str]]): The input text to be embedded.
+        
+    Returns:
+        torch.Tensor: Embedding of the input text.
+    """
+    if isinstance(text, str):   
+        text = [text]
+
+    if 'GritLM' in model_name:
+        emb = get_gritlm_embeddings(text, model_name, **encode_kwargs)
+    elif 'LLM2Vec' in model_name:
+        emb = get_llm2vec_embeddings(text, model_name, **encode_kwargs)
+    elif 'all-mpnet-base-v2' in model_name or 'dunzhang/stella_en_1.5B_v5' in model_name:
+        emb = get_sentence_transformer_embeddings(text, model_name, **encode_kwargs)
+    else:
+        emb = get_api_embeddings(text, model_name, **encode_kwargs)
+    return emb.view(len(text), -1)
+
+
+if __name__ == '__main__':
+    print("Testing ModelForSTaRKQA...")

Reasoning/utils.py

@@ -0,0 +1,116 @@
+from Reasoning.text_retrievers.bm25 import BM25
+from Reasoning.text_retrievers.ada import Ada
+from Reasoning.text_retrievers.contriever import Contriever
+
+
+def combine_dicts(dicts_list, pred_dict):
+    if len(dicts_list) == 1:
+        return dicts_list[0]
+    combined_dict = {}
+    
+    for d in dicts_list:
+        for key, value in d.items():
+            if key in combined_dict:
+                # for route dict, the values are lists, keep the longest list
+                if len(value) > len(combined_dict[key]):
+                    combined_dict[key] = value
+            else:
+                combined_dict[key] = value
+
+    # if the two reasoning paths have intersection, only keep the keys in pred_dict 
+    combined_dict = {key: combined_dict[key] for key in pred_dict.keys()}
+    
+    
+    return combined_dict
+    
+def fix_length(paths_dict):
+    max_length = 3
+    new_paths_dict = {}
+    
+    for key, value in paths_dict.items():
+        if len(value) > max_length:
+            value = value[-max_length:]
+        if len(value) < max_length:
+            # padding with -1 at the beginning
+            value = [-1] * (max_length - len(value)) + value
+        new_paths_dict[key] = value
+        
+    return new_paths_dict
+            
+    
+
+def parse_metapath(metapath):
+    """
+        input: metapath: "paper -> author -> paper <- paper"
+        output: routes: [['paper', 'author', 'paper'], ['paper', 'paper']]
+    """
+    
+    def parse(remain_list, direction):
+        """
+            input: remain_list: ["paper", "->", "author", "->", "paper", "<-", "paper"]
+                direction: "->"
+            output: route: ["paper", "author", "paper"]
+                    remain_list: ["paper", "<-", "paper"]
+        """
+        route = []
+        i = 0
+        while i < len(remain_list)-1 and remain_list[i+1] == direction:
+            route.append(remain_list[i])
+            i += 2
+        route.append(remain_list[i])
+        
+        if direction == "<-":
+            route.reverse()    
+            
+        remain_list = None if len(remain_list) == i+1 else remain_list[i:]
+            
+        return route, remain_list
+    
+    
+    remain_list = metapath.split(' ')
+    # print(f"111, {remain_list}")
+    
+    if len(remain_list) == 1: # single node
+        return [remain_list]
+    
+    routes = []
+    while remain_list is not None:
+        if remain_list[1] == "<-":
+            route, remain_list = parse(remain_list, "<-")
+        
+        elif remain_list[1] == "->":
+            route, remain_list = parse(remain_list, "->")
+            
+        else: 
+            # raise ValueError(f"Invalid metapath: {metapath}")
+            return None
+        
+        routes.append(route)
+            
+    return routes 
+
+
+def get_text_retriever(dataset_name, retriever_name, skb, **kwargs):
+    if retriever_name == "bm25":
+        return BM25(skb, dataset_name)
+    elif retriever_name == "ada":
+        return Ada(skb, dataset_name, kwargs.get("device", 'cuda'))
+    elif retriever_name == "contriever":
+        return Contriever(skb, dataset_name, kwargs.get("device", 'cuda'))
+    else:
+        raise ValueError(f"Invalid retriever name: {retriever_name}")
+
+
+def get_scorer(dataset_name, scorer_name, skb, **kwargs):
+    if scorer_name == "bm25":
+        return BM25(skb, dataset_name)
+    elif scorer_name == "ada":
+        return Ada(skb, dataset_name, kwargs.get("device",'cuda'))
+    elif scorer_name == "contriever":
+        return Contriever(skb, dataset_name, kwargs.get("device", 'cuda'))
+    else:
+        raise ValueError(f"Invalid scorer name: {scorer_name}")
+    
+
+if __name__ == "__main__":
+    print(f"Test utils")

Reranking/data/checkpoints/amazon/best.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a4ef56bf1095b28283a5bc451556e238b31d36285df39b1a84b9c5c65f1900a2
+size 804607

Reranking/data/checkpoints/mag/best.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b25c04fcb5fa1210fe8ca10bb7c1b3f3aad592ec0e501a1ccc330b92772550f1
+size 804620

Reranking/data/checkpoints/prime/best.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d1e4e24531067e798b975e33689cc87645a2daa9d736df1e8ac0474f72154ce0
+size 804633

Reranking/rerank.py

@@ -0,0 +1,375 @@
+import sys
+import os
+sys.path.append(os.path.dirname(os.path.abspath(os.getcwd())))
+from stark_qa import load_skb
+
+
+from torch.utils.data import Dataset, DataLoader
+import torch
+from tqdm import tqdm
+import numpy as np
+import torch.nn as nn
+
+from Reranking.utils import move_to_cuda, seed_everything
+from Reranking.rerankers.path import PathReranker
+import torch.nn.functional as F
+import argparse
+import pickle as pkl
+
+
+
+
+class TestDataset(Dataset):
+    """        
+        data format: {
+            "query": query,
+            "pred_dict": {node_id: score},
+            'score_vector_dict': {node_id: [bm25, bm_25, bm25, ada]},
+            "text_emb_dict": {node_id: text_emb},
+            "ans_ids": [],
+        }
+
+"""
+
+    def __init__(self, saved_data, args):
+        
+        print(f"Start processing test dataset...")
+        self.text2emb_dict = saved_data['text2emb_dict']
+        self.data = saved_data['data']  
+        
+        self.text_emb_matrix = list(self.text2emb_dict.values())
+        self.text_emb_matrix = torch.concat(self.text_emb_matrix, dim=0)
+        
+        # make the mapping between the key of text2emb_dict and the index of text_emb_matrix
+        self.text2idx_dict = {key: idx for idx, key in enumerate(self.text2emb_dict.keys())}
+        
+        self.args = args
+        
+             
+        
+      
+    def __len__(self):
+        return len(self.data) 
+
+    def __getitem__(self, idx):
+        
+        if self.args.dataset_name == 'amazon':
+            # change from the str to index
+            self.data[idx]['text_emb_dict'] = {key: self.text2idx_dict[value] for key, value in self.data[idx]['text_emb_dict'].items()}
+        else:
+            # sort the pred_dict by the score
+            pred_dict = self.data[idx]['pred_dict']
+            sorted_ids = sorted(pred_dict.keys(), key=lambda x: pred_dict[x], reverse=True)
+            # get the top 50 candidates
+            sorted_ids = sorted_ids[:50]
+            # get the score vector
+            self.data[idx]['score_vector_dict'] = {key: self.data[idx]['score_vector_dict'][key] for key in sorted_ids}
+            # get the symb_enc_dict
+            self.data[idx]['symb_enc_dict'] = {key: self.data[idx]['symb_enc_dict'][key] for key in sorted_ids}
+            # change from the str to index
+            self.data[idx]['text_emb_dict'] = {key: self.text2idx_dict[value] for key, value in self.data[idx]['text_emb_dict'].items()}
+            self.data[idx]['text_emb_dict'] = {key: self.data[idx]['text_emb_dict'][key] for key in sorted_ids}
+    
+                
+
+        return self.data[idx]
+    
+    
+    def collate_batch(self, batch):
+        
+        # q
+        batch_q = [batch[i]['query'] for i in range(len(batch))]
+        q_text = batch_q
+        
+        # c
+        batch_c = [list(batch[i]['score_vector_dict'].keys()) for i in range(len(batch))] # [batch, 100]
+        batch_c = torch.tensor(batch_c)
+        c_score_vector = [list(batch[i]['score_vector_dict'].values()) for i in range(len(batch))] # [batch, 100, 4]
+        c_score_vector = torch.tensor(c_score_vector)
+        c_score_vector = c_score_vector[:, :, :self.args.vector_dim]
+        
+        # c_symb_enc
+        c_symb_enc = [list(batch[i]['symb_enc_dict'].values()) for i in range(len(batch))]
+        c_symb_enc = torch.tensor(c_symb_enc) # [bs, 100, 3]
+    
+        # c_text_emb
+        c_text_emb = [self.text_emb_matrix[list(batch[i]['text_emb_dict'].values())].unsqueeze(0) for i in range(len(batch))]
+        c_text_emb = torch.concat(c_text_emb, dim=0) # [bs, 100, 768]
+        
+        
+        # ans_ids
+        ans_ids = [batch[i]['ans_ids'] for i in range(len(batch))] # list of ans_ids
+        
+        # pred_ids
+        pred_ids = batch_c.tolist()
+            
+        
+        # Create a dictionary for the batch
+        feed_dict = {
+            'query': q_text,
+            'c_score_vector': c_score_vector,
+            'c_text_emb': c_text_emb,
+            'c_symb_enc': c_symb_enc,
+            'ans_ids': ans_ids,
+            'pred_ids': pred_ids
+
+        }
+        
+        
+        return feed_dict
+    
+    
+#  ***** batch_evaluator *****
+def batch_evaluator(skb, scores_cand, ans_ids, batch):
+
+    results = {}
+    
+    # **** batch wise evaluation ****
+    # evaluate
+    candidates_ids = skb.candidate_ids
+    id_to_idx = {candidate_id: idx for idx, candidate_id in enumerate(candidates_ids)}
+    
+    
+    # initialize the pred_matrix
+    pred_matrix = torch.zeros((scores_cand.shape[0],len(candidates_ids)))
+    
+    
+    # get the index of each pred_ids
+    # flatten the pred_ids
+    flat_pred_ids = torch.tensor(batch['pred_ids']).flatten().tolist()
+    
+    
+    # get the index of each pred_ids
+    pred_idx = [id_to_idx[pred_id] for pred_id in flat_pred_ids]
+    
+    
+    # reshape the pred_idx
+    pred_idx = torch.tensor(pred_idx).reshape(scores_cand.shape[0], -1) # [bs, 100]
+    
+    # move pred_matrix to the device
+    pred_matrix = pred_matrix.to(scores_cand.device)
+    
+    # advanced indexing
+    pred_matrix[torch.arange(scores_cand.shape[0]).unsqueeze(1), pred_idx] = scores_cand.squeeze(-1) # [bs, num_candidates]
+    
+    
+    # Create a mapping from candidate IDs to their indices for faster lookup
+    
+
+    # Flatten ans_ids to a single list and map them to indices
+    flat_ans_idx = [id_to_idx[a_id] for sublist in ans_ids for a_id in sublist]
+
+    # Create the row indices for ans_matrix corresponding to the answers
+    row_indices = torch.repeat_interleave(torch.arange(len(ans_ids)), torch.tensor([len(sublist) for sublist in ans_ids]))
+
+    # Create the answer matrix
+    ans_matrix = torch.zeros((scores_cand.shape[0], len(candidates_ids)), device=scores_cand.device)
+    ans_matrix[row_indices, torch.tensor(flat_ans_idx, device=scores_cand.device)] = 1
+
+
+    
+    # batch computing hit1
+    # find the index of the max score
+    max_score, max_idx = torch.max(pred_matrix, dim=1)
+    # check the label of the max idx
+    batch_hit1 = ans_matrix[torch.arange(scores_cand.shape[0]), max_idx]
+    hit1_list = batch_hit1.tolist()
+    
+    
+    # batch computing hit@5
+    _, top5_idx = torch.topk(pred_matrix, 5, dim=1)
+    batch_hit5 = ans_matrix[torch.arange(scores_cand.shape[0]).unsqueeze(1), top5_idx]
+    
+    # max with each row
+    batch_hit5 = torch.max(batch_hit5, dim=1)[0]
+    hit5_list = batch_hit5.tolist()
+    
+    
+    
+    # batch computing recall@20
+    _, top20_idx = torch.topk(pred_matrix, 20, dim=1)
+    batch_recall20 = ans_matrix[torch.arange(scores_cand.shape[0]).unsqueeze(1), top20_idx]
+    # sum with each row
+    batch_recall20 = torch.sum(batch_recall20, dim=1)
+    # divide by the sum of the ans_matrix along the row
+    batch_recall20 = batch_recall20 / torch.sum(ans_matrix, dim=1)
+    recall20_list = batch_recall20.tolist()
+    
+    
+    
+    # batch computing mrr
+    # find the highest rank of the answer
+    _, rank_idx = torch.sort(pred_matrix, dim=1, descending=True)
+    # query the answer matrix with the rank_idx
+    batch_mrr = ans_matrix[torch.arange(scores_cand.shape[0]).unsqueeze(1), rank_idx]
+    # find the first rank of the answer
+    batch_mrr = torch.argmax(batch_mrr, dim=1)
+    # add 1 to the rank
+    batch_mrr += 1
+    # divide by the rank
+    batch_mrr = 1 / batch_mrr.float()
+    mrr_list = batch_mrr.tolist()
+    
+
+    results['hit@1'] = hit1_list
+    results['hit@5'] = hit5_list
+    results['recall@20'] = recall20_list
+    results['mrr'] = mrr_list
+    
+
+    return results
+    
+
+
+# ***** evaluate *****
+@torch.no_grad()
+def evaluate(router, test_loader, skb):
+
+    
+    router.eval()
+
+    all_results = {
+        "hit@1": [],
+        "hit@5": [],
+        "recall@20": [],
+        "mrr": []
+    }
+    avg_results = {
+        "hit@1": 0,
+        "hit@5": 0,
+        "recall@20": 0,
+        "mrr": 0
+    }
+    
+    
+    # save the scores and ans_ids, and pred_ids
+    pred_list = []
+    scores_cand_list = []
+    ans_ids_list = []
+    print(f"Start evaluating...")
+    # use tqdm to show the progress
+    for idx, batch in enumerate(tqdm(test_loader, desc='Evaluating', position=0)):
+        # print(f"idx: {idx}")
+        batch = move_to_cuda(batch)
+        
+        # Check if the model is wrapped in DataParallel
+        if isinstance(router, nn.DataParallel):
+            scores_cand = router.module.eval_batch(batch) # q_emb: [bs, 100], c_emb: [bs*100, 100]
+        else:
+            scores_cand = router.eval_batch(batch)
+
+
+        # ans_ids
+        ans_ids = batch['ans_ids']
+           
+        results = batch_evaluator(skb, scores_cand, ans_ids, batch)
+                
+        
+        for key in results.keys():
+            all_results[key].extend(results[key])
+            
+        # save the scores and ans_ids, and pred_ids
+        pred_list.extend(batch['pred_ids'])
+        scores_cand_list.extend(scores_cand.cpu().tolist())
+        ans_ids_list.extend(ans_ids)
+        
+    
+    
+    for key in avg_results.keys():
+        avg_results[key] = np.mean(all_results[key])
+    
+    print(f"Results: {avg_results}")
+    
+
+    
+    return avg_results
+
+
+def parse_args():
+    
+    parser = argparse.ArgumentParser(description="Run PathRouter with dynamic combinations of embeddings.")
+    
+    # dataset_name
+    parser.add_argument("--dataset_name", type=str, default="mag", help="Name of the dataset.")
+    
+    # Add arguments for model configurations
+    parser.add_argument("--device", type=str, default="cuda", help="Device to run the model (e.g., 'cuda' or 'cpu').")
+
+    
+    # add concat_num
+    parser.add_argument("--concat_num", type=int, default=0, help="Number of concatenation of embeddings.")
+    
+    # checkpoint save path
+    parser.add_argument("--checkpoint_path", type=str, default="./data/checkpoints", help="Path saves the checkpoints.")
+    
+    # similarity vector dim
+    parser.add_argument("--vector_dim", type=int, default=4, help="Dimension of the similarity vector.")
+    
+    
+    # Parse the base arguments
+    args = parser.parse_args()
+    return args
+
+
+def get_concat_num(combo):
+    """
+    Determine the value of concat_num based on the combination of embeddings.
+    - score_vec adds +1
+    - text_emb adds +1
+    - symb_enc adds +3
+    """
+    concat_num = 0
+    if combo.get("score_vec", False):  # If score_vec is True
+        concat_num += 1
+    if combo.get("text_emb", False):  # If text_emb is True
+        concat_num += 1
+    if combo.get("symb_enc", False):  # If symb_enc is True
+        concat_num += 3
+        
+        
+    return concat_num
+
+
+def run(test_data, skb, dataset_name):
+    
+    
+    
+    test_size = 64
+    test_dataset = TestDataset(test_data, args=args)
+    test_loader = DataLoader(test_dataset, batch_size=test_size, num_workers=32, collate_fn=test_dataset.collate_batch)
+    
+    # load the model
+    print(f"Load the model...")
+    args.checkpoint_path = args.checkpoint_path + f"/{dataset_name}/best.pth"
+    router = PathReranker(socre_vector_input_dim=4, text_emb_input_dim=768, symb_enc_dim=3, args=args)
+    checkpoint = torch.load(args.checkpoint_path)
+    router.load_state_dict(checkpoint)
+    router = router.to(args.device)
+    
+    # evalute
+    test_results = evaluate(router, test_loader, skb)
+    print(f"Test evaluation")
+    print(test_results)
+    
+    return test_results
+
+if __name__ == "__main__":
+    
+    combo = {
+        "text_emb": True,
+        "score_vec": True,
+        "symb_enc": True
+    }
+    concat_num = get_concat_num(combo)
+    
+    base_args = parse_args()
+    args = argparse.Namespace(**vars(base_args), **combo)
+    args.concat_num = concat_num
+    dataset_name = args.dataset_name
+    
+    test_data_path = f"../{dataset_name}_test.pkl"
+    with open(test_data_path, 'rb') as f:
+        test_data = pkl.load(f)
+    skb = load_skb(dataset_name)
+    results = run(test_data, skb, dataset_name)
+    

Reranking/rerankers/node.py

@@ -0,0 +1,87 @@
+import torch.nn as nn
+
+    
+
+# ***** Reranking Model *****
+class Contriever(nn.Module):
+    def __init__(self, encoder, emb_dim=768):
+        super(Contriever, self).__init__()
+        self.encoder = encoder
+        self.emb_dim = emb_dim
+    
+    def mean_pooling(self, token_embs, mask):
+        token_embs = token_embs.masked_fill(~mask[..., None].bool(), 0.0)
+        sentence_embeddings = token_embs.sum(dim=1) / (mask.sum(dim=1)[..., None].clamp(min=1e-9))
+        return sentence_embeddings
+    
+    def encode_seq(self, input_ids, attention_mask, token_type_ids=None):
+        # Combine inputs into a dictionary
+        enc = {'input_ids': input_ids, 'attention_mask': attention_mask}
+        if token_type_ids is not None:
+            enc['token_type_ids'] = token_type_ids
+
+        outputs = self.encoder(**enc)
+        # Mean pooling of last hidden states
+        embedded = self.mean_pooling(outputs[0], attention_mask)
+        # print(f"777, {embedded.shape}")
+        return embedded
+
+    def get_text_emb(self, input_ids, attention_mask, token_type_ids):
+        emb = self.encode_seq(input_ids=input_ids, attention_mask=attention_mask, token_type_ids=token_type_ids)
+        
+        return emb
+    
+    def eval_batch(self, batch):
+        # q_emb: [batch_size/num_gpus, token_dim]
+        q_emb = self.get_text_emb(batch['q_enc_input_ids'], batch['q_enc_attention_mask'], batch['q_enc_token_type_ids'])
+        
+        # c_emb: [batch_size * num_candidates/num_gpus, token_dim]
+        c_emb = self.get_text_emb(batch['c_enc_input_ids'], batch['c_enc_attention_mask'], batch['c_enc_token_type_ids'])
+        
+
+        
+        return q_emb, c_emb
+    
+    def forward(self, batch):
+        
+        # q_emb: [batch_size/num_gpus, token_dim]
+        q_emb = self.get_text_emb(batch['q_enc_input_ids'], batch['q_enc_attention_mask'], batch['q_enc_token_type_ids'])
+        # p_emb: [batch_size*max_len/num_gpus, token_dim]
+        p_emb = self.get_text_emb(batch['pos_enc_input_ids'], batch['pos_enc_attention_mask'], batch['pos_enc_token_type_ids'])
+        # n_emb: [batch_size*max_len*num_sampled_negs/num_gpus, token_dim]
+        n_emb = self.get_text_emb(batch['neg_enc_input_ids'], batch['neg_enc_attention_mask'], batch['neg_enc_token_type_ids'])
+        
+
+        
+        return q_emb, p_emb, n_emb
+
+
+# ***** Reranking Model *****
+class NodeRouter(nn.Module):
+    def __init__(self, input_dim=2, output_dim=1, emb_dim=128):
+        super(NodeRouter, self).__init__()
+        self.fc1 = nn.Linear(input_dim, emb_dim)
+        self.fc2 = nn.Linear(emb_dim, output_dim)
+        self.relu = nn.ReLU()
+    
+    def eval_batch(self, batch):
+        scores_cand = self.fc1(batch['c_scores'])   
+        scores_cand = self.relu(scores_cand)
+        scores_cand = self.fc2(scores_cand)
+        scores_cand = self.relu(scores_cand)
+        
+        return scores_cand
+    
+    def forward(self, batch):
+        scores_pos = self.fc1(batch['p_scores'])
+        scores_neg = self.fc1(batch['n_scores'])
+        scores_pos = self.relu(scores_pos)
+        scores_neg = self.relu(scores_neg)
+        
+        scores_pos = self.fc2(scores_pos)
+        scores_neg = self.fc2(scores_neg)
+        scores_pos = self.relu(scores_pos)
+        scores_neg = self.relu(scores_neg)
+        
+        return scores_pos, scores_neg
+        

Reranking/rerankers/path.py

@@ -0,0 +1,107 @@
+import torch.nn as nn
+import torch
+
+# ***** Reranking Model *****
+# two linear layers
+class PathReranker(nn.Module):
+    def __init__(self, socre_vector_input_dim=4, text_emb_input_dim=768, symb_enc_dim=3, output_dim=1, emb_dim=256, args=None):
+        super(PathReranker, self).__init__()
+        self.score_vec_enc = nn.Linear(socre_vector_input_dim, emb_dim)
+        self.text_emb_enc = nn.Linear(text_emb_input_dim, emb_dim)
+        self.symb_enc = nn.Embedding(symb_enc_dim, emb_dim)
+        self.fc1 = nn.Linear(emb_dim*args.concat_num, output_dim)
+        self.fc2 = nn.Linear(output_dim, 1)
+        self.relu = nn.ReLU()
+        self.args = args
+        
+        
+    
+    def eval_batch(self, batch):
+        
+        embeddings = []
+        
+        if self.args.text_emb:
+            # Encode the text embedding and apply ReLU
+            text_emb_c = self.relu(self.text_emb_enc(batch['c_text_emb'])) # [bs, 100, emb_dim]
+            embeddings.append(text_emb_c)
+            
+        
+        if self.args.score_vec:
+            # Encode the score vector and apply ReLU
+           score_vector_c = self.relu(self.score_vec_enc(batch['c_score_vector'])) # [bs, 100, emb_dim]
+        
+           embeddings.append(score_vector_c)
+        
+        
+        if self.args.symb_enc:
+            # encode the symbolic embedding and apply ReLU
+            symb_enc_c = self.relu(self.symb_enc(batch['c_symb_enc']))
+            # reshape the symbolic embedding
+            symb_enc_c = torch.reshape(symb_enc_c, (symb_enc_c.shape[0], symb_enc_c.shape[1], -1))
+            embeddings.append(symb_enc_c)
+            
+            
+        if len(embeddings) > 1:
+            emb_c = torch.cat(embeddings, dim=-1)
+        else:
+            emb_c = embeddings[0]
+            
+        
+        # Feed the concatenated embeddings to the final layer
+        emb_c = self.fc1(emb_c) # [bs, 100, emb_dim]
+        scores_c = self.fc2(emb_c) # [bs, 100, 1]
+        
+        return scores_c
+
+    
+    def forward(self, batch):    
+        
+        embeddings_pos = []
+        embeddings_neg = []
+        
+        if self.args.text_emb:
+            # Encode the text embedding and apply ReLU
+            text_emb_pos = self.relu(self.text_emb_enc(batch['p_text_emb']))
+            text_emb_neg = self.relu(self.text_emb_enc(batch['n_text_emb']))
+            embeddings_pos.append(text_emb_pos)
+            embeddings_neg.append(text_emb_neg)
+        
+        if self.args.score_vec:
+            # Encode the score vector and apply ReLU
+            score_vector_pos = self.relu(self.score_vec_enc(batch['p_score_vector']))
+            score_vector_neg = self.relu(self.score_vec_enc(batch['n_score_vector']))
+            embeddings_pos.append(score_vector_pos)
+            embeddings_neg.append(score_vector_neg)
+        
+        
+        if self.args.symb_enc:
+            # encode the symbolic embedding and apply ReLU
+            symb_enc_pos = self.relu(self.symb_enc(batch['p_symb_enc']))
+            # reshape the symbolic embedding
+            symb_enc_pos = torch.reshape(symb_enc_pos, (symb_enc_pos.shape[0], -1))
+            
+            symb_enc_neg = self.relu(self.symb_enc(batch['n_symb_enc']))
+            # reshape the symbolic embedding
+            symb_enc_neg = torch.reshape(symb_enc_neg, (symb_enc_neg.shape[0], symb_enc_neg.shape[1], -1)) # [bs, neg_sp, path_len * emb_dim]
+            
+            embeddings_pos.append(symb_enc_pos)
+            embeddings_neg.append(symb_enc_neg)
+        
+        
+        if len(embeddings_pos) > 1:
+            pos = torch.cat(embeddings_pos, dim=-1)
+            neg = torch.cat(embeddings_neg, dim=-1)
+        else:
+            pos = embeddings_pos[0]
+            neg = embeddings_neg[0]
+                
+        
+        
+        # Feed the concatenated embeddings to the final layer
+        pos = self.fc1(pos)
+        neg = self.fc1(neg)
+        scores_pos = self.fc2(pos)
+        scores_neg = self.fc2(neg)
+        
+        return scores_pos, scores_neg
+    

Reranking/train_eval_path_amazon.py

@@ -0,0 +1,694 @@
+import sys
+sys.path.append(os.path.dirname(os.path.abspath(os.getcwd())))
+
+import pickle as pkl
+from torch.utils.data import Dataset, DataLoader
+import torch
+from tqdm import tqdm
+import wandb
+import numpy as np
+import time
+from torch_scatter import segment_csr, scatter_mean
+from itertools import product
+from transformers import AutoModel, AutoTokenizer, AutoConfig
+import torch.nn as nn
+from torch.nn import CrossEntropyLoss
+import random
+from collections import defaultdict
+import os
+
+from Reranking.utils import move_to_cuda, seed_everything
+from Reranking.rerankers.path import PathReranker
+from utils import ModelForSTaRKQA
+from stark_qa import load_qa, load_skb
+import torch.nn.functional as F
+import argparse
+import json
+import time
+
+
+seed_everything(42)
+
+# ***** Dataset *****
+class TrainDataset(Dataset):
+    """
+    Custom Dataset for the training data.
+    Each instance contains multiple positive and negative candidates.
+    """
+    def __init__(self, saved_data, max_neg_candidates=100):
+        """
+            10s for 1000 data 
+        """
+        print(f"start processing training dataset...")
+        s_time = time.time()
+        self.max_neg_candidates = max_neg_candidates        
+        self.sorted_query2neg = defaultdict(list)
+        
+
+        self.text2emb_dict = saved_data['text2emb_dict']
+        self.data = saved_data['data']  
+
+        
+        # separage neg and pos, and prepare query, pos pairs
+        new_data = []
+        
+        for i in range(len(self.data)):
+            neg_ids = []
+            pos_ids = []
+            item = self.data[i]
+            
+            
+            candidates_dict = item['pred_dict']
+            ans_ids = item['ans_ids']
+            # pos_ids = ans_ids
+            for ans_id in ans_ids:
+                if ans_id in candidates_dict.keys():
+                    pos_ids.append(ans_id)
+            neg_ids = list(set(candidates_dict.keys()) - set(pos_ids))
+            
+            # load scores vector
+            score_vector_dict = item['score_vector_dict']
+            
+            # load the text path, str format
+            text_emb_dict = item['text_emb_dict']
+            
+            # load the symb_enc_dict
+            symb_enc_dict = item['symb_enc_dict']
+
+            
+            self.data[i]['pos_ids'] = pos_ids
+            self.data[i]['neg_ids'] = neg_ids
+            
+            query = item['query']
+            for pos_id in pos_ids:
+                new_data.append((query, score_vector_dict[pos_id], self.text2emb_dict[text_emb_dict[pos_id]], symb_enc_dict[pos_id]))
+                
+            
+            # print(f"new_data: {new_data}")
+
+            neg_dict = {neg_id: candidates_dict[neg_id] for neg_id in neg_ids}
+            sorted_neg_ids = sorted(neg_dict.keys(), key=lambda x: neg_dict[x], reverse=True) # return list
+            
+                
+            self.sorted_query2neg[query] = [(score_vector_dict[neg_id], self.text2emb_dict[text_emb_dict[neg_id]], symb_enc_dict[neg_id]) for neg_id in sorted_neg_ids]        
+                
+            
+        self.data = new_data
+        print(f"Complete data preparation")
+        print(f"Time: {time.time() - s_time}")
+        
+        
+        
+
+    def __len__(self):
+        return len(self.data) 
+
+    def __getitem__(self, idx):
+
+        return self.data[idx]
+        
+    def collate_batch(self, pairs):
+        s_time = time.time()
+        
+        # q
+        batch_q = [pair[0] for pair in pairs] # q is text
+        q_text = batch_q
+        # print(f"q111, {q_text}")
+        
+        
+        # pos
+        # get the score vector
+        batch_p_score_vector = [pair[1] for pair in pairs] # p is score vector
+        batch_p_score_vector = torch.tensor(batch_p_score_vector) # [bs, 4]
+        batch_p_score_vector = batch_p_score_vector[:, :args.vector_dim]
+        # get the text emb
+        batch_p_text_emb = [pair[2] for pair in pairs] # p is text emb
+        batch_p_text_emb = torch.concat(batch_p_text_emb, dim=0) # [bs, 768]
+        # get the symb_enc
+        batch_p_symb_enc = [pair[3] for pair in pairs] # p is symb_enc
+        batch_p_symb_enc = torch.tensor(batch_p_symb_enc) # [bs, 3]
+        
+        
+        # Negative samples
+        batch_n = [random.choices(self.sorted_query2neg[query], k=self.max_neg_candidates) for query in batch_q] # allow duplicates
+        
+        
+        # get the score vector
+        batch_n_score_vector = [pair[0] for sublist in batch_n for pair in sublist]
+        batch_n_score_vector = torch.tensor(batch_n_score_vector) # [bs*100, 4]
+        # reshape to [bs, 100, 4]
+        batch_n_score_vector = batch_n_score_vector.reshape(len(batch_q), self.max_neg_candidates, -1) # [bs, 100, 4]
+        batch_n_score_vector = batch_n_score_vector[:, :, :args.vector_dim]
+        
+        # get the text emb
+        batch_n_text_emb = [pair[1] for sublist in batch_n for pair in sublist] 
+        batch_n_text_emb = torch.concat(batch_n_text_emb, dim=0) # [bs*100, 768]
+        # reshape to [bs, 100, 768]
+        batch_n_text_emb = batch_n_text_emb.reshape(len(batch_q), self.max_neg_candidates, -1) # [bs, 100, 768]
+        
+        # get the symb_enc
+        batch_n_symb_enc = [pair[2] for sublist in batch_n for pair in sublist]
+        batch_n_symb_enc = torch.tensor(batch_n_symb_enc) # [bs*100, 3]
+        # reshape to [bs, 100, 3]
+        batch_n_symb_enc = batch_n_symb_enc.reshape(len(batch_q), self.max_neg_candidates, -1) # [bs, 100, 3]
+            
+        
+        
+        
+        
+        # Create a dictionary for the batch
+        feed_dict = {
+            'query': q_text,
+            'p_score_vector': batch_p_score_vector,
+            'p_text_emb': batch_p_text_emb,
+            'p_symb_enc': batch_p_symb_enc,
+            'n_score_vector': batch_n_score_vector,
+            'n_text_emb': batch_n_text_emb,
+            'n_symb_enc': batch_n_symb_enc,
+            
+        }
+
+
+        return feed_dict
+    
+
+class TestDataset(Dataset):
+    """        
+        data format: {
+            "query": query,
+            "pred_dict": {node_id: score},
+            'score_vector_dict': {node_id: [bm25, bm_25, bm25, ada]},
+            "text_emb_dict": {node_id: text_emb},
+            "ans_ids": [],
+        }
+
+"""
+
+    def __init__(self, saved_data):
+        
+        print(f"Start processing test dataset...")
+        self.text2emb_dict = saved_data['text2emb_dict']
+        self.data = saved_data['data']  
+        
+        self.text_emb_matrix = list(self.text2emb_dict.values())
+        self.text_emb_matrix = torch.concat(self.text_emb_matrix, dim=0)
+        
+        # make the mapping between the key of text2emb_dict and the index of text_emb_matrix
+        self.text2idx_dict = {key: idx for idx, key in enumerate(self.text2emb_dict.keys())}
+        
+        print(f"Complete data preparation: {len(self.data)}")
+        
+             
+        
+      
+    def __len__(self):
+        return len(self.data) 
+
+    def __getitem__(self, idx):
+        # ***** amazon *****
+        # change from the str to index
+        self.data[idx]['text_emb_dict'] = {key: self.text2idx_dict[value] for key, value in self.data[idx]['text_emb_dict'].items()}
+        
+        return self.data[idx]
+    
+    def collate_batch(self, batch):
+        
+        # q
+        batch_q = [batch[i]['query'] for i in range(len(batch))]
+        q_text = batch_q
+        
+        # c
+        batch_c = [list(batch[i]['score_vector_dict'].keys()) for i in range(len(batch))] # [batch, 100]
+        batch_c = torch.tensor(batch_c)
+        # print(f"111, {batch_c.shape}")
+        c_score_vector = [list(batch[i]['score_vector_dict'].values()) for i in range(len(batch))] # [batch, 100, 4]
+        c_score_vector = torch.tensor(c_score_vector)[:, :, :args.vector_dim] # [batch, 100, 4]
+        
+        
+        # print(f"222, {c_vector.shape}")
+        # c_text_emb
+        # c_text_emb = [torch.concat(list(batch[i]['text_emb_dict'].values()), dim=0).unsqueeze(0) for i in range(len(batch))]
+        c_text_emb = [self.text_emb_matrix[list(batch[i]['text_emb_dict'].values())].unsqueeze(0) for i in range(len(batch))]
+        c_text_emb = torch.concat(c_text_emb, dim=0) # [bs, 100, 768]
+        
+        # c_symb_enc
+        c_symb_enc = [list(batch[i]['symb_enc_dict'].values()) for i in range(len(batch))]
+        c_symb_enc = torch.tensor(c_symb_enc) # [bs, 100, 3]
+        
+        
+        # ans_ids
+        ans_ids = [batch[i]['ans_ids'] for i in range(len(batch))] # list of ans_ids
+        
+        # pred_ids
+        pred_ids = batch_c.tolist()
+            
+        
+        # Create a dictionary for the batch
+        feed_dict = {
+            'query': q_text,
+            'c_score_vector': c_score_vector,
+            'c_text_emb': c_text_emb,
+            'c_symb_enc': c_symb_enc,
+            'ans_ids': ans_ids,
+            'pred_ids': pred_ids
+
+        }
+
+        
+        
+        return feed_dict
+    
+
+# ******* loss function ********
+def loss_fn(scores_pos, scores_neg):
+    
+    loss_fct = CrossEntropyLoss(ignore_index=-1)
+
+    # Combine scores
+    scores = torch.cat([scores_pos, scores_neg.squeeze(-1)], dim=1)  # B x (1 + max_neg_candidates*B)
+    # print(f"scores: {scores.shape}")
+    
+    # Create target
+    target = torch.zeros(scores.size(0), dtype=torch.long).to(scores.device)
+
+    # Compute loss
+    loss = loss_fct(scores, target)
+    
+    return loss
+
+# ***** pairwise loss *****
+def pairwise_loss(scores_pos, scores_neg, margin=0.5):
+    # scores_pos: [bs, 1]
+    # scores_neg: [bs, 100, 1]
+
+    # Compute loss
+    differences = scores_pos.unsqueeze(1) - scores_neg - margin # [bs, 100, 1]
+    differences = differences.view(-1)  # [bs*100]
+    loss = F.relu(-differences).mean()  # Standard pairwise loss
+
+    return loss
+
+
+def batch_evaluator(skb, scores_cand, ans_ids, batch):
+
+    results = {}
+    
+    # **** batch wise evaluation ****
+    # evaluate
+    candidates_ids = skb.candidate_ids
+    id_to_idx = {candidate_id: idx for idx, candidate_id in enumerate(candidates_ids)}
+    
+    
+    # initialize the pred_matrix
+    pred_matrix = torch.zeros((scores_cand.shape[0],len(candidates_ids)))
+    
+    
+    # get the index of each pred_ids
+    # flatten the pred_ids
+    flat_pred_ids = torch.tensor(batch['pred_ids']).flatten().tolist()
+    
+    
+    # get the index of each pred_ids
+    pred_idx = [id_to_idx[pred_id] for pred_id in flat_pred_ids]
+    
+    
+    # reshape the pred_idx
+    pred_idx = torch.tensor(pred_idx).reshape(scores_cand.shape[0], -1) # [bs, 100]
+    
+    # move pred_matrix to the device
+    pred_matrix = pred_matrix.to(scores_cand.device)
+    
+    # advanced indexing
+    pred_matrix[torch.arange(scores_cand.shape[0]).unsqueeze(1), pred_idx] = scores_cand.squeeze(-1) # [bs, num_candidates]
+    
+
+    
+    # Create a mapping from candidate IDs to their indices for faster lookup
+    
+
+    # Flatten ans_ids to a single list and map them to indices
+    flat_ans_idx = [id_to_idx[a_id] for sublist in ans_ids for a_id in sublist]
+
+    # Create the row indices for ans_matrix corresponding to the answers
+    row_indices = torch.repeat_interleave(torch.arange(len(ans_ids)), torch.tensor([len(sublist) for sublist in ans_ids]))
+
+    # Create the answer matrix
+    ans_matrix = torch.zeros((scores_cand.shape[0], len(candidates_ids)), device=scores_cand.device)
+    ans_matrix[row_indices, torch.tensor(flat_ans_idx, device=scores_cand.device)] = 1
+
+
+    
+    
+    
+    # batch computing hit1
+    # find the index of the max score
+    max_score, max_idx = torch.max(pred_matrix, dim=1)
+    # check the label of the max idx
+    batch_hit1 = ans_matrix[torch.arange(scores_cand.shape[0]), max_idx]
+    hit1_list = batch_hit1.tolist()
+    
+    
+
+    
+    # batch computing hit@5
+    _, top5_idx = torch.topk(pred_matrix, 5, dim=1)
+    batch_hit5 = ans_matrix[torch.arange(scores_cand.shape[0]).unsqueeze(1), top5_idx]
+    
+    # max with each row
+    batch_hit5 = torch.max(batch_hit5, dim=1)[0]
+    hit5_list = batch_hit5.tolist()
+    
+    
+    
+    # batch computing recall@20
+    _, top20_idx = torch.topk(pred_matrix, 20, dim=1)
+    batch_recall20 = ans_matrix[torch.arange(scores_cand.shape[0]).unsqueeze(1), top20_idx]
+    # sum with each row
+    batch_recall20 = torch.sum(batch_recall20, dim=1)
+    # divide by the sum of the ans_matrix along the row
+    batch_recall20 = batch_recall20 / torch.sum(ans_matrix, dim=1)
+    recall20_list = batch_recall20.tolist()
+    
+    
+    
+    
+    # batch computing mrr
+    # find the highest rank of the answer
+    _, rank_idx = torch.sort(pred_matrix, dim=1, descending=True)
+    # query the answer matrix with the rank_idx
+    batch_mrr = ans_matrix[torch.arange(scores_cand.shape[0]).unsqueeze(1), rank_idx]
+    # find the first rank of the answer
+    batch_mrr = torch.argmax(batch_mrr, dim=1)
+    # add 1 to the rank
+    batch_mrr += 1
+    # divide by the rank
+    batch_mrr = 1 / batch_mrr.float()
+    mrr_list = batch_mrr.tolist()
+    
+    
+    
+
+    results['hit@1'] = hit1_list
+    results['hit@5'] = hit5_list
+    results['recall@20'] = recall20_list
+    results['mrr'] = mrr_list
+    
+
+    
+
+    
+    
+    return results
+    
+
+
+# ***** evaluate *****
+@torch.no_grad()
+def evaluate(reranker, test_loader):
+
+    
+    reranker.eval()
+
+    all_results = {
+        "hit@1": [],
+        "hit@5": [],
+        "recall@20": [],
+        "mrr": []
+    }
+    avg_results = {
+        "hit@1": 0,
+        "hit@5": 0,
+        "recall@20": 0,
+        "mrr": 0
+    }
+    
+    
+    # save the scores and ans_ids, and pred_ids
+    pred_list = []
+    scores_cand_list = []
+    ans_ids_list = []
+    # use tqdm to show the progress
+    for idx, batch in enumerate(tqdm(test_loader, desc='Evaluating', position=0)):
+        batch = move_to_cuda(batch)
+        
+        # Check if the model is wrapped in DataParallel
+        if isinstance(reranker, nn.DataParallel):
+            scores_cand = reranker.module.eval_batch(batch) # q_emb: [bs, 100], c_emb: [bs*100, 100]
+        else:
+            scores_cand = reranker.eval_batch(batch)
+
+
+        # ans_ids
+        ans_ids = batch['ans_ids']
+           
+        results = batch_evaluator(skb, scores_cand, ans_ids, batch)
+                
+        
+        for key in results.keys():
+            all_results[key].extend(results[key])
+            
+        # save the scores and ans_ids, and pred_ids
+        pred_list.extend(batch['pred_ids'])
+        scores_cand_list.extend(scores_cand.cpu().tolist())
+        ans_ids_list.extend(ans_ids)
+        
+    
+    
+    for key in avg_results.keys():
+        avg_results[key] = np.mean(all_results[key])
+    
+    print(f"Results: {avg_results}")
+    
+
+    
+    return avg_results
+    
+    
+# ***** train *****
+def main(train_data, val_data, test_data, skb, dataset_name, args):
+
+
+    epochs = args.epochs
+    device = args.device
+
+    train_size = args.train_batch_size
+    test_size = 64
+
+    train_dataset = TrainDataset(train_data)
+    train_loader = DataLoader(train_dataset, batch_size=train_size, num_workers=32, collate_fn=train_dataset.collate_batch, drop_last=True)
+
+    test_dataset = TestDataset(test_data)
+    test_loader = DataLoader(test_dataset, batch_size=test_size, num_workers=32, collate_fn=test_dataset.collate_batch)
+    
+    val_dataset = TestDataset(val_data)
+    val_loader = DataLoader(val_dataset, batch_size=test_size, num_workers=32, collate_fn=val_dataset.collate_batch)
+    
+
+    # ***** Model *****
+    reranker = PathReranker(socre_vector_input_dim=args.vector_dim, text_emb_input_dim=768, symb_enc_dim=3, args=args)
+    save_dir = f"./data/checkpoints/{dataset_name}/path"
+    os.makedirs(save_dir, exist_ok=True)
+
+    reranker.to(device)
+    # # parallel processing
+    reranker = nn.DataParallel(reranker)
+
+
+    optimizer = torch.optim.Adam(reranker.parameters(), lr=args.lr)
+    best_val_hit1 = float('-inf')
+
+    
+    val_results = evaluate(reranker, val_loader)
+    print(f"Val evaluation")
+    print(val_results)
+    
+    
+    test_results = evaluate(reranker, test_loader)
+    print(f"Test evaluation")
+    print(test_results)
+    
+    # log both val and test results
+    wandb.log({'val_mrr': val_results['mrr'], 'val_hit1': val_results['hit@1'], 'val_hit5': val_results['hit@5'], 'val_recall@20': val_results['recall@20'],
+               'test_mrr': test_results['mrr'], 'test_hit1': test_results['hit@1'], 'test_hit5': test_results['hit@5'], 'test_recall@20': test_results['recall@20']})
+    
+    best_test_results = {}
+    for epoch in tqdm(range(epochs), desc='Training Epochs', position=0):
+        total_loss = 0.0
+        reranker.train()
+        count = 0
+        total_instances = 0
+
+        for batch in tqdm(train_loader):
+            # print(batch)
+            batch = move_to_cuda(batch)
+            # print(batch)
+            
+            scores_pos, scores_neg = reranker(batch)
+            
+            # batch_loss = pairwise_loss(scores_pos, scores_neg)
+            batch_loss = loss_fn(scores_pos, scores_neg)
+            
+            # clear optimizer 
+            optimizer.zero_grad()
+            batch_loss.backward()
+            optimizer.step()
+            
+            # total_loss += batch_loss.item()
+            count += 1 
+            # compute the average loss
+            total_instances += scores_pos.shape[0]
+            total_loss += batch_loss.item()
+                        
+            
+        train_loss = total_loss / total_instances
+        
+        print(f"Epoch {epoch+1}/{epochs}, Average Train Loss: {train_loss}")
+        
+        
+        
+        val_results = evaluate(reranker, val_loader)
+        print(f"Val evaluation")
+        print(val_results)
+        
+        
+        test_results = evaluate(reranker, test_loader)
+        print(f"Test evaluation")
+        print(test_results)
+        
+        # log both val and test results
+        wandb.log({'val_mrr': val_results['mrr'], 'val_hit1': val_results['hit@1'], 'val_hit5': val_results['hit@5'], 'val_recall@20': val_results['recall@20'],
+                'test_mrr': test_results['mrr'], 'test_hit1': test_results['hit@1'], 'test_hit5': test_results['hit@5'], 'test_recall@20': test_results['recall@20'],
+                'train_loss': train_loss})
+        
+        
+        # save the best model when val hit1 is the highest
+        hit1 = val_results['hit@1']
+        if best_val_hit1 < hit1:
+            best_val_hit1 = hit1
+            
+            save_path = f"{save_dir}/best_{best_val_hit1}.pth"
+            
+            if isinstance(reranker, nn.DataParallel):
+                torch.save(reranker.module.state_dict(), save_path)
+            else:
+                torch.save(reranker.state_dict(), save_path)
+            print(f"Checkpoint saved at epoch {epoch+1} with test hits@1 {hit1}")
+            
+            args.checkpoint_path = save_path
+            best_test_results = test_results
+        
+        
+                
+    # save last epoch checkopint
+    save_path = f"{save_dir}/last_{hit1}.pth"
+    if isinstance(reranker, nn.DataParallel):
+        torch.save(reranker.module.state_dict(), save_path)
+    else:
+        torch.save(reranker.state_dict(), save_path)
+    print(f"Final checkpoint saved at {save_path}")
+    
+    
+    
+    # ***** save the results *****
+    results = []
+    results.append(
+        {
+            "config": vars(args),
+            "test_results": best_test_results
+        }
+    )
+    # save the results to json
+    timestamp = time.strftime("%Y%m%d-%H%M%S")
+    output_dir = f"./data/outputs/{dataset_name}"
+    os.makedirs(output_dir, exist_ok=True)
+    with open(f"{output_dir}/results_{timestamp}.json", "w") as f:
+        json.dump(results, f, indent=4)
+        
+    print(best_test_results)
+    
+    
+def get_concat_num(combo):
+    """
+    Determine the value of concat_num based on the combination of embeddings.
+    - score_vec adds +1
+    - text_emb adds +1
+    - symb_enc adds +3
+    """
+    concat_num = 0
+    if combo.get("score_vec", False):  # If score_vec is True
+        concat_num += 1
+    if combo.get("text_emb", False):  # If text_emb is True
+        concat_num += 1
+    if combo.get("symb_enc", False):  # If symb_enc is True
+        concat_num += 3
+        
+        
+    return concat_num
+    
+    
+    
+def parse_args():
+    
+    parser = argparse.ArgumentParser(description="Run Pathreranker with dynamic combinations of embeddings.")
+    
+    # Add arguments for model configurations
+    parser.add_argument("--train_batch_size", type=int, default=64, help="Batch size for training or evaluation.")
+    parser.add_argument("--lr", type=float, default=1e-4, help="Learning rate for optimizer.")
+    parser.add_argument("--epochs", type=int, default=100, help="Number of epochs to train the model.")
+    parser.add_argument("--device", type=str, default="cuda", help="Device to run the model (e.g., 'cuda' or 'cpu').")
+    
+    # Add arguments for the dataset
+    parser.add_argument("--dataset_name", type=str, default="amazon", help="Name of the dataset to use.")
+    # paths
+    parser.add_argument("--train_path", type=str, default=f"../amazon_train.pkl", help="Path to the training data.")
+    parser.add_argument("--test_path", type=str, default=f"../amazon_test.pkl", help="Path to the test data.")
+    parser.add_argument("--val_path", type=str, default=f"../amazon_val.pkl", help="Path to the validation data.")
+    
+    # add concat_num
+    parser.add_argument("--concat_num", type=int, default=0, help="Number of concatenation of embeddings.")
+    
+    # checkpoint save path
+    parser.add_argument("--checkpoint_path", type=str, default="", help="Path to save the checkpoints.")
+    parser.add_argument("--vector_dim", type=int, default=4, help="Dimension of the similarity vector.")
+    
+    
+    # Parse the base arguments
+    args = parser.parse_args()
+    return args
+    
+    
+if __name__ == "__main__":
+    
+    base_args = parse_args()
+    test_path = base_args.test_path
+    train_path = base_args.train_path
+    val_path = base_args.val_path
+    dataset_name = base_args.dataset_name
+    
+    with open(test_path, "rb") as f:
+        test_data = pkl.load(f)
+        
+    with open(train_path, "rb") as f:
+        train_data = pkl.load(f)
+        
+    with open(val_path, "rb") as f:
+        val_data = pkl.load(f)
+    
+    # load skb
+    skb = load_skb(dataset_name)
+    
+    # set all
+    combo = {
+        "text_emb": True,
+        "score_vec": True,
+        "symb_enc": True
+    }
+    concat_num = get_concat_num(combo)
+    
+    wandb.init(project=f'reranking-{dataset_name}', name=f"path")
+    args = argparse.Namespace(**vars(base_args), **combo)
+    args.concat_num = concat_num
+    
+    
+    main(train_data, val_data, test_data, skb, dataset_name, args)
+        

Reranking/train_eval_path_mag.py

@@ -0,0 +1,694 @@
+import sys
+sys.path.append(os.path.dirname(os.path.abspath(os.getcwd())))
+
+import pickle as pkl
+from torch.utils.data import Dataset, DataLoader
+import torch
+from tqdm import tqdm
+import wandb
+import numpy as np
+import time
+from torch_scatter import segment_csr, scatter_mean
+from itertools import product
+from transformers import AutoModel, AutoTokenizer, AutoConfig
+import torch.nn as nn
+from torch.nn import CrossEntropyLoss
+import random
+from collections import defaultdict
+import os
+
+from Reranking.utils import move_to_cuda, seed_everything
+from Reranking.rerankers.path import PathReranker
+from utils import ModelForSTaRKQA
+from stark_qa import load_qa, load_skb
+import torch.nn.functional as F
+import argparse
+import json
+import time
+
+seed_everything(42)
+
+# ***** Dataset *****
+class TrainDataset(Dataset):
+    """
+    Custom Dataset for the training data.
+    Each instance contains multiple positive and negative candidates.
+    """
+    def __init__(self, saved_data, max_neg_candidates=100):
+        """
+            10s for 1000 data 
+        """
+        print(f"start processing training dataset...")
+        s_time = time.time()
+        self.max_neg_candidates = max_neg_candidates        
+        self.sorted_query2neg = defaultdict(list)
+        
+
+        self.text2emb_dict = saved_data['text2emb_dict']
+        self.data = saved_data['data']  
+
+        
+        # separage neg and pos, and prepare query, pos pairs
+        new_data = []
+        
+        for i in range(len(self.data)):
+            neg_ids = []
+            pos_ids = []
+            item = self.data[i]
+            
+            
+            candidates_dict = item['pred_dict']
+            ans_ids = item['ans_ids']
+            # pos_ids = ans_ids
+            for ans_id in ans_ids:
+                if ans_id in candidates_dict.keys():
+                    pos_ids.append(ans_id)
+            neg_ids = list(set(candidates_dict.keys()) - set(pos_ids))
+            
+            # load scores vector
+            score_vector_dict = item['score_vector_dict']
+            
+            # load the text path, str format
+            text_emb_dict = item['text_emb_dict']
+            
+            # load the symb_enc_dict
+            symb_enc_dict = item['symb_enc_dict']
+
+            
+            self.data[i]['pos_ids'] = pos_ids
+            self.data[i]['neg_ids'] = neg_ids
+            
+            query = item['query']
+            for pos_id in pos_ids:
+                new_data.append((query, score_vector_dict[pos_id], self.text2emb_dict[text_emb_dict[pos_id]], symb_enc_dict[pos_id]))
+                
+            
+            # print(f"new_data: {new_data}")
+
+            neg_dict = {neg_id: candidates_dict[neg_id] for neg_id in neg_ids}
+            sorted_neg_ids = sorted(neg_dict.keys(), key=lambda x: neg_dict[x], reverse=True) # return list
+            
+                
+            self.sorted_query2neg[query] = [(score_vector_dict[neg_id], self.text2emb_dict[text_emb_dict[neg_id]], symb_enc_dict[neg_id]) for neg_id in sorted_neg_ids]        
+                
+            
+        self.data = new_data
+        print(f"Complete data preparation")
+        print(f"Time: {time.time() - s_time}")
+        
+        
+        
+
+    def __len__(self):
+        return len(self.data) 
+
+    def __getitem__(self, idx):
+
+        return self.data[idx]
+        
+    def collate_batch(self, pairs):
+        s_time = time.time()
+        
+        # q
+        batch_q = [pair[0] for pair in pairs] # q is text
+        q_text = batch_q
+        # print(f"q111, {q_text}")
+        
+        
+        # pos
+        # get the score vector
+        batch_p_score_vector = [pair[1] for pair in pairs] # p is score vector
+        batch_p_score_vector = torch.tensor(batch_p_score_vector) # [bs, 4]
+        batch_p_score_vector = batch_p_score_vector[:, :args.vector_dim]
+        # get the text emb
+        batch_p_text_emb = [pair[2] for pair in pairs] # p is text emb
+        batch_p_text_emb = torch.concat(batch_p_text_emb, dim=0) # [bs, 768]
+        # get the symb_enc
+        batch_p_symb_enc = [pair[3] for pair in pairs] # p is symb_enc
+        batch_p_symb_enc = torch.tensor(batch_p_symb_enc) # [bs, 3]
+        
+        
+        # Negative samples
+        batch_n = [random.choices(self.sorted_query2neg[query], k=self.max_neg_candidates) for query in batch_q] # allow duplicates
+        
+        
+        # get the score vector
+        batch_n_score_vector = [pair[0] for sublist in batch_n for pair in sublist]
+        batch_n_score_vector = torch.tensor(batch_n_score_vector) # [bs*100, 4]
+        # reshape to [bs, 100, 4]
+        batch_n_score_vector = batch_n_score_vector.reshape(len(batch_q), self.max_neg_candidates, -1) # [bs, 100, 4]
+        batch_n_score_vector = batch_n_score_vector[:, :, :args.vector_dim]
+        
+        # get the text emb
+        batch_n_text_emb = [pair[1] for sublist in batch_n for pair in sublist] 
+        batch_n_text_emb = torch.concat(batch_n_text_emb, dim=0) # [bs*100, 768]
+        # reshape to [bs, 100, 768]
+        batch_n_text_emb = batch_n_text_emb.reshape(len(batch_q), self.max_neg_candidates, -1) # [bs, 100, 768]
+        
+        # get the symb_enc
+        batch_n_symb_enc = [pair[2] for sublist in batch_n for pair in sublist]
+        batch_n_symb_enc = torch.tensor(batch_n_symb_enc) # [bs*100, 3]
+        # reshape to [bs, 100, 3]
+        batch_n_symb_enc = batch_n_symb_enc.reshape(len(batch_q), self.max_neg_candidates, -1) # [bs, 100, 3]
+            
+        
+        
+        
+        
+        # Create a dictionary for the batch
+        feed_dict = {
+            'query': q_text,
+            'p_score_vector': batch_p_score_vector,
+            'p_text_emb': batch_p_text_emb,
+            'p_symb_enc': batch_p_symb_enc,
+            'n_score_vector': batch_n_score_vector,
+            'n_text_emb': batch_n_text_emb,
+            'n_symb_enc': batch_n_symb_enc,
+            
+        }
+
+
+        return feed_dict
+    
+
+class TestDataset(Dataset):
+    """        
+        data format: {
+            "query": query,
+            "pred_dict": {node_id: score},
+            'score_vector_dict': {node_id: [bm25, bm_25, bm25, ada]},
+            "text_emb_dict": {node_id: text_emb},
+            "ans_ids": [],
+        }
+
+"""
+
+    def __init__(self, saved_data):
+        
+        print(f"Start processing test dataset...")
+        self.text2emb_dict = saved_data['text2emb_dict']
+        self.data = saved_data['data']  
+        
+        self.text_emb_matrix = list(self.text2emb_dict.values())
+        self.text_emb_matrix = torch.concat(self.text_emb_matrix, dim=0)
+        
+        # make the mapping between the key of text2emb_dict and the index of text_emb_matrix
+        self.text2idx_dict = {key: idx for idx, key in enumerate(self.text2emb_dict.keys())}
+        
+        print(f"Complete data preparation: {len(self.data)}")
+        
+             
+        
+      
+    def __len__(self):
+        return len(self.data) 
+
+    def __getitem__(self, idx):
+
+        
+        # sort the pred_dict by the score
+        pred_dict = self.data[idx]['pred_dict']
+        sorted_ids = sorted(pred_dict.keys(), key=lambda x: pred_dict[x], reverse=True)
+        # get the top 50 candidates
+        sorted_ids = sorted_ids[:50]
+        # get the score vector
+        self.data[idx]['score_vector_dict'] = {key: self.data[idx]['score_vector_dict'][key] for key in sorted_ids}
+        # get the symb_enc_dict
+        self.data[idx]['symb_enc_dict'] = {key: self.data[idx]['symb_enc_dict'][key] for key in sorted_ids}
+        # change from the str to index
+        self.data[idx]['text_emb_dict'] = {key: self.text2idx_dict[value] for key, value in self.data[idx]['text_emb_dict'].items()}
+        self.data[idx]['text_emb_dict'] = {key: self.data[idx]['text_emb_dict'][key] for key in sorted_ids}
+        
+        return self.data[idx]
+    
+    def collate_batch(self, batch):
+        
+        # q
+        batch_q = [batch[i]['query'] for i in range(len(batch))]
+        q_text = batch_q
+        
+        # c
+        batch_c = [list(batch[i]['score_vector_dict'].keys()) for i in range(len(batch))] # [batch, 100]
+        batch_c = torch.tensor(batch_c)
+        # print(f"111, {batch_c.shape}")
+        c_score_vector = [list(batch[i]['score_vector_dict'].values()) for i in range(len(batch))] # [batch, 100, 4]
+        c_score_vector = torch.tensor(c_score_vector)[:, :, :args.vector_dim] # [batch, 100, 4]
+        
+        
+        # print(f"222, {c_vector.shape}")
+        # c_text_emb
+        # c_text_emb = [torch.concat(list(batch[i]['text_emb_dict'].values()), dim=0).unsqueeze(0) for i in range(len(batch))]
+        c_text_emb = [self.text_emb_matrix[list(batch[i]['text_emb_dict'].values())].unsqueeze(0) for i in range(len(batch))]
+        c_text_emb = torch.concat(c_text_emb, dim=0) # [bs, 100, 768]
+        
+        # c_symb_enc
+        c_symb_enc = [list(batch[i]['symb_enc_dict'].values()) for i in range(len(batch))]
+        c_symb_enc = torch.tensor(c_symb_enc) # [bs, 100, 3]
+        
+        
+        # ans_ids
+        ans_ids = [batch[i]['ans_ids'] for i in range(len(batch))] # list of ans_ids
+        
+        # pred_ids
+        pred_ids = batch_c.tolist()
+            
+        
+        # Create a dictionary for the batch
+        feed_dict = {
+            'query': q_text,
+            'c_score_vector': c_score_vector,
+            'c_text_emb': c_text_emb,
+            'c_symb_enc': c_symb_enc,
+            'ans_ids': ans_ids,
+            'pred_ids': pred_ids
+
+        }
+
+        
+        
+        return feed_dict
+    
+
+# ******* loss function ********
+def loss_fn(scores_pos, scores_neg):
+    
+    loss_fct = CrossEntropyLoss(ignore_index=-1)
+
+    # Combine scores
+    scores = torch.cat([scores_pos, scores_neg.squeeze(-1)], dim=1)  # B x (1 + max_neg_candidates*B)
+    # print(f"scores: {scores.shape}")
+    
+    # Create target
+    target = torch.zeros(scores.size(0), dtype=torch.long).to(scores.device)
+
+    # Compute loss
+    loss = loss_fct(scores, target)
+    
+    return loss
+
+
+
+def batch_evaluator(skb, scores_cand, ans_ids, batch):
+
+    results = {}
+    
+    # **** batch wise evaluation ****
+    # evaluate
+    candidates_ids = skb.candidate_ids
+    id_to_idx = {candidate_id: idx for idx, candidate_id in enumerate(candidates_ids)}
+    
+    
+    # initialize the pred_matrix
+    pred_matrix = torch.zeros((scores_cand.shape[0],len(candidates_ids)))
+    
+    
+    # get the index of each pred_ids
+    # flatten the pred_ids
+    flat_pred_ids = torch.tensor(batch['pred_ids']).flatten().tolist()
+    
+    
+    # get the index of each pred_ids
+    pred_idx = [id_to_idx[pred_id] for pred_id in flat_pred_ids]
+    
+    
+    # reshape the pred_idx
+    pred_idx = torch.tensor(pred_idx).reshape(scores_cand.shape[0], -1) # [bs, 100]
+    
+    # move pred_matrix to the device
+    pred_matrix = pred_matrix.to(scores_cand.device)
+    
+    # advanced indexing
+    pred_matrix[torch.arange(scores_cand.shape[0]).unsqueeze(1), pred_idx] = scores_cand.squeeze(-1) # [bs, num_candidates]
+    
+
+    
+    # Create a mapping from candidate IDs to their indices for faster lookup
+    
+
+    # Flatten ans_ids to a single list and map them to indices
+    flat_ans_idx = [id_to_idx[a_id] for sublist in ans_ids for a_id in sublist]
+
+    # Create the row indices for ans_matrix corresponding to the answers
+    row_indices = torch.repeat_interleave(torch.arange(len(ans_ids)), torch.tensor([len(sublist) for sublist in ans_ids]))
+
+    # Create the answer matrix
+    ans_matrix = torch.zeros((scores_cand.shape[0], len(candidates_ids)), device=scores_cand.device)
+    ans_matrix[row_indices, torch.tensor(flat_ans_idx, device=scores_cand.device)] = 1
+
+
+    
+    
+    
+    # batch computing hit1
+    # find the index of the max score
+    max_score, max_idx = torch.max(pred_matrix, dim=1)
+    # check the label of the max idx
+    batch_hit1 = ans_matrix[torch.arange(scores_cand.shape[0]), max_idx]
+    hit1_list = batch_hit1.tolist()
+    
+    
+
+    
+    # batch computing hit@5
+    _, top5_idx = torch.topk(pred_matrix, 5, dim=1)
+    batch_hit5 = ans_matrix[torch.arange(scores_cand.shape[0]).unsqueeze(1), top5_idx]
+    
+    # max with each row
+    batch_hit5 = torch.max(batch_hit5, dim=1)[0]
+    hit5_list = batch_hit5.tolist()
+    
+    
+    
+    # batch computing recall@20
+    _, top20_idx = torch.topk(pred_matrix, 20, dim=1)
+    batch_recall20 = ans_matrix[torch.arange(scores_cand.shape[0]).unsqueeze(1), top20_idx]
+    # sum with each row
+    batch_recall20 = torch.sum(batch_recall20, dim=1)
+    # divide by the sum of the ans_matrix along the row
+    batch_recall20 = batch_recall20 / torch.sum(ans_matrix, dim=1)
+    recall20_list = batch_recall20.tolist()
+    
+    
+    
+    
+    # batch computing mrr
+    # find the highest rank of the answer
+    _, rank_idx = torch.sort(pred_matrix, dim=1, descending=True)
+    # query the answer matrix with the rank_idx
+    batch_mrr = ans_matrix[torch.arange(scores_cand.shape[0]).unsqueeze(1), rank_idx]
+    # find the first rank of the answer
+    batch_mrr = torch.argmax(batch_mrr, dim=1)
+    # add 1 to the rank
+    batch_mrr += 1
+    # divide by the rank
+    batch_mrr = 1 / batch_mrr.float()
+    mrr_list = batch_mrr.tolist()
+    
+    
+    
+
+    results['hit@1'] = hit1_list
+    results['hit@5'] = hit5_list
+    results['recall@20'] = recall20_list
+    results['mrr'] = mrr_list
+    
+
+    
+
+    
+    
+    return results
+    
+
+
+# ***** evaluate *****
+@torch.no_grad()
+def evaluate(reranker, test_loader):
+
+    
+    reranker.eval()
+
+    all_results = {
+        "hit@1": [],
+        "hit@5": [],
+        "recall@20": [],
+        "mrr": []
+    }
+    avg_results = {
+        "hit@1": 0,
+        "hit@5": 0,
+        "recall@20": 0,
+        "mrr": 0
+    }
+    
+    
+    # save the scores and ans_ids, and pred_ids
+    pred_list = []
+    scores_cand_list = []
+    ans_ids_list = []
+    # use tqdm to show the progress
+    for idx, batch in enumerate(tqdm(test_loader, desc='Evaluating', position=0)):
+        batch = move_to_cuda(batch)
+        
+        # Check if the model is wrapped in DataParallel
+        if isinstance(reranker, nn.DataParallel):
+            scores_cand = reranker.module.eval_batch(batch) # q_emb: [bs, 100], c_emb: [bs*100, 100]
+        else:
+            scores_cand = reranker.eval_batch(batch)
+
+
+        # ans_ids
+        ans_ids = batch['ans_ids']
+           
+        results = batch_evaluator(skb, scores_cand, ans_ids, batch)
+                
+        
+        for key in results.keys():
+            all_results[key].extend(results[key])
+            
+        # save the scores and ans_ids, and pred_ids
+        pred_list.extend(batch['pred_ids'])
+        scores_cand_list.extend(scores_cand.cpu().tolist())
+        ans_ids_list.extend(ans_ids)
+        
+    
+    
+    for key in avg_results.keys():
+        avg_results[key] = np.mean(all_results[key])
+    
+    print(f"Results: {avg_results}")
+    
+
+    
+    return avg_results
+    
+    
+# ***** train *****
+def main(train_data, val_data, test_data, skb, dataset_name, args):
+
+
+    epochs = args.epochs
+    device = args.device
+
+    train_size = args.train_batch_size
+    test_size = 64
+
+    train_dataset = TrainDataset(train_data)
+    train_loader = DataLoader(train_dataset, batch_size=train_size, num_workers=32, collate_fn=train_dataset.collate_batch, drop_last=True)
+
+    test_dataset = TestDataset(test_data)
+    test_loader = DataLoader(test_dataset, batch_size=test_size, num_workers=32, collate_fn=test_dataset.collate_batch)
+    
+    val_dataset = TestDataset(val_data)
+    val_loader = DataLoader(val_dataset, batch_size=test_size, num_workers=32, collate_fn=val_dataset.collate_batch)
+    
+
+    # ***** Model *****
+    reranker = PathReranker(socre_vector_input_dim=args.vector_dim, text_emb_input_dim=768, symb_enc_dim=3, args=args)
+    save_dir = f"./data/checkpoints/{dataset_name}/path"
+    os.makedirs(save_dir, exist_ok=True)
+
+    reranker.to(device)
+    # # parallel processing
+    reranker = nn.DataParallel(reranker)
+
+
+    optimizer = torch.optim.Adam(reranker.parameters(), lr=args.lr)
+    best_val_hit1 = float('-inf')
+
+    
+    val_results = evaluate(reranker, val_loader)
+    print(f"Val evaluation")
+    print(val_results)
+    
+    
+    test_results = evaluate(reranker, test_loader)
+    print(f"Test evaluation")
+    print(test_results)
+    
+    # log both val and test results
+    wandb.log({'val_mrr': val_results['mrr'], 'val_hit1': val_results['hit@1'], 'val_hit5': val_results['hit@5'], 'val_recall@20': val_results['recall@20'],
+               'test_mrr': test_results['mrr'], 'test_hit1': test_results['hit@1'], 'test_hit5': test_results['hit@5'], 'test_recall@20': test_results['recall@20']})
+    
+    best_test_results = {}
+    for epoch in tqdm(range(epochs), desc='Training Epochs', position=0):
+        total_loss = 0.0
+        reranker.train()
+        count = 0
+        total_instances = 0
+
+        for batch in tqdm(train_loader):
+            # print(batch)
+            batch = move_to_cuda(batch)
+            # print(batch)
+            
+            scores_pos, scores_neg = reranker(batch)
+            
+            # batch_loss = pairwise_loss(scores_pos, scores_neg)
+            batch_loss = loss_fn(scores_pos, scores_neg)
+            
+            # clear optimizer 
+            optimizer.zero_grad()
+            batch_loss.backward()
+            optimizer.step()
+            
+            # total_loss += batch_loss.item()
+            count += 1 
+            # compute the average loss
+            total_instances += scores_pos.shape[0]
+            total_loss += batch_loss.item()
+                        
+            
+        train_loss = total_loss / total_instances
+        
+        print(f"Epoch {epoch+1}/{epochs}, Average Train Loss: {train_loss}")
+        
+        
+        
+        val_results = evaluate(reranker, val_loader)
+        print(f"Val evaluation")
+        print(val_results)
+        
+        
+        test_results = evaluate(reranker, test_loader)
+        print(f"Test evaluation")
+        print(test_results)
+        
+        # log both val and test results
+        wandb.log({'val_mrr': val_results['mrr'], 'val_hit1': val_results['hit@1'], 'val_hit5': val_results['hit@5'], 'val_recall@20': val_results['recall@20'],
+                'test_mrr': test_results['mrr'], 'test_hit1': test_results['hit@1'], 'test_hit5': test_results['hit@5'], 'test_recall@20': test_results['recall@20'],
+                'train_loss': train_loss})
+        
+        
+        # save the best model when val hit1 is the highest
+        hit1 = val_results['hit@1']
+        if best_val_hit1 < hit1:
+            best_val_hit1 = hit1
+            
+            save_path = f"{save_dir}/best_{best_val_hit1}.pth"
+            
+            if isinstance(reranker, nn.DataParallel):
+                torch.save(reranker.module.state_dict(), save_path)
+            else:
+                torch.save(reranker.state_dict(), save_path)
+            print(f"Checkpoint saved at epoch {epoch+1} with test hits@1 {hit1}")
+            
+            args.checkpoint_path = save_path
+            best_test_results = test_results
+        
+        
+                
+    # save last epoch checkopint
+    save_path = f"{save_dir}/last_{hit1}.pth"
+    if isinstance(reranker, nn.DataParallel):
+        torch.save(reranker.module.state_dict(), save_path)
+    else:
+        torch.save(reranker.state_dict(), save_path)
+    print(f"Final checkpoint saved at {save_path}")
+    
+    
+    
+    # ***** save the results *****
+    results = []
+    results.append(
+        {
+            "config": vars(args),
+            "test_results": best_test_results
+        }
+    )
+    # save the results to json
+    timestamp = time.strftime("%Y%m%d-%H%M%S")
+    output_dir = f"./data/outputs/{dataset_name}"
+    os.makedirs(output_dir, exist_ok=True)
+    with open(f"{output_dir}/results_{timestamp}.json", "w") as f:
+        json.dump(results, f, indent=4)
+    
+    print(best_test_results)
+        
+        
+    
+def get_concat_num(combo):
+    """
+    Determine the value of concat_num based on the combination of embeddings.
+    - score_vec adds +1
+    - text_emb adds +1
+    - symb_enc adds +3
+    """
+    concat_num = 0
+    if combo.get("score_vec", False):  # If score_vec is True
+        concat_num += 1
+    if combo.get("text_emb", False):  # If text_emb is True
+        concat_num += 1
+    if combo.get("symb_enc", False):  # If symb_enc is True
+        concat_num += 3
+        
+        
+    return concat_num
+    
+    
+    
+def parse_args():
+    
+    parser = argparse.ArgumentParser(description="Run Pathreranker with dynamic combinations of embeddings.")
+    
+    # Add arguments for model configurations
+    parser.add_argument("--train_batch_size", type=int, default=5, help="Batch size for training or evaluation.")
+    parser.add_argument("--lr", type=float, default=1e-4, help="Learning rate for optimizer.")
+    parser.add_argument("--epochs", type=int, default=100, help="Number of epochs to train the model.")
+    parser.add_argument("--device", type=str, default="cuda", help="Device to run the model (e.g., 'cuda' or 'cpu').")
+    
+    # Add arguments for the dataset
+    parser.add_argument("--dataset_name", type=str, default="mag", help="Name of the dataset to use.")
+    # paths
+    parser.add_argument("--train_path", type=str, default=f"../mag_train.pkl", help="Path to the training data.")
+    parser.add_argument("--test_path", type=str, default=f"../mag_test.pkl", help="Path to the test data.")
+    parser.add_argument("--val_path", type=str, default=f"../mag_val.pkl", help="Path to the validation data.")
+    
+    # add concat_num
+    parser.add_argument("--concat_num", type=int, default=0, help="Number of concatenation of embeddings.")
+    
+    # checkpoint save path
+    parser.add_argument("--checkpoint_path", type=str, default="", help="Path to save the checkpoints.")
+    parser.add_argument("--vector_dim", type=int, default=4, help="Dimension of the similarity vector.")
+    
+    
+    # Parse the base arguments
+    args = parser.parse_args()
+    return args
+    
+    
+if __name__ == "__main__":
+    
+    base_args = parse_args()
+
+    test_path = base_args.test_path
+    train_path = base_args.train_path
+    val_path = base_args.val_path
+    dataset_name = base_args.dataset_name
+    
+    with open(test_path, "rb") as f:
+        test_data = pkl.load(f)
+        
+    with open(train_path, "rb") as f:
+        train_data = pkl.load(f)
+        
+    with open(val_path, "rb") as f:
+        val_data = pkl.load(f)
+    
+    # load skb
+    skb = load_skb(dataset_name)
+    
+    # set all
+    combo = {
+        "text_emb": True,
+        "score_vec": True,
+        "symb_enc": True
+    }
+    concat_num = get_concat_num(combo)
+    
+    wandb.init(project=f'Reranking-{dataset_name}', name=f"path")
+    args = argparse.Namespace(**vars(base_args), **combo)
+    args.concat_num = concat_num
+    
+    main(train_data, val_data, test_data, skb, dataset_name, args)
+        

Reranking/train_eval_path_prime.py

@@ -0,0 +1,691 @@
+import sys
+import os
+sys.path.append(os.path.dirname(os.path.abspath(os.getcwd())))
+
+import pickle as pkl
+from torch.utils.data import Dataset, DataLoader
+import torch
+from tqdm import tqdm
+import wandb
+import numpy as np
+import time
+import torch.nn as nn
+from torch.nn import CrossEntropyLoss
+import random
+from collections import defaultdict
+
+
+from Reranking.utils import move_to_cuda, seed_everything
+from Reranking.rerankers.path import PathReranker
+from stark_qa import load_qa, load_skb
+import torch.nn.functional as F
+import argparse
+import json
+import time
+
+# set the seed
+seed_everything(42)
+
+# ***** Dataset *****
+class TrainDataset(Dataset):
+    """
+    Custom Dataset for the training data.
+    Each instance contains multiple positive and negative candidates.
+    """
+    def __init__(self, saved_data, max_neg_candidates=100):
+        """
+            10s for 1000 data 
+        """
+        print(f"start processing training dataset...")
+        s_time = time.time()
+        self.max_neg_candidates = max_neg_candidates        
+        self.sorted_query2neg = defaultdict(list)
+        
+
+        self.text2emb_dict = saved_data['text2emb_dict']
+        self.data = saved_data['data']  
+
+        
+        # separage neg and pos, and prepare query, pos pairs
+        new_data = []
+        
+        for i in range(len(self.data)):
+            neg_ids = []
+            pos_ids = []
+            item = self.data[i]
+            
+            
+            candidates_dict = item['pred_dict']
+            ans_ids = item['ans_ids']
+            # pos_ids = ans_ids
+            for ans_id in ans_ids:
+                if ans_id in candidates_dict.keys():
+                    pos_ids.append(ans_id)
+            neg_ids = list(set(candidates_dict.keys()) - set(pos_ids))
+            
+            # load scores vector
+            score_vector_dict = item['score_vector_dict']
+            
+            # load the text path, str format
+            text_emb_dict = item['text_emb_dict']
+            
+            # load the symb_enc_dict
+            symb_enc_dict = item['symb_enc_dict']
+
+            
+            self.data[i]['pos_ids'] = pos_ids
+            self.data[i]['neg_ids'] = neg_ids
+            
+            query = item['query']
+            for pos_id in pos_ids:
+                new_data.append((query, score_vector_dict[pos_id], self.text2emb_dict[text_emb_dict[pos_id]], symb_enc_dict[pos_id]))
+                
+            
+            # print(f"new_data: {new_data}")
+
+            neg_dict = {neg_id: candidates_dict[neg_id] for neg_id in neg_ids}
+            sorted_neg_ids = sorted(neg_dict.keys(), key=lambda x: neg_dict[x], reverse=True) # return list
+            
+                
+            self.sorted_query2neg[query] = [(score_vector_dict[neg_id], self.text2emb_dict[text_emb_dict[neg_id]], symb_enc_dict[neg_id]) for neg_id in sorted_neg_ids]        
+                
+            
+        self.data = new_data
+        print(f"Complete data preparation")
+        print(f"Time: {time.time() - s_time}")
+        
+        
+        
+
+    def __len__(self):
+        return len(self.data) 
+
+    def __getitem__(self, idx):
+
+        return self.data[idx]
+        
+    def collate_batch(self, pairs):
+        s_time = time.time()
+        
+        # q
+        batch_q = [pair[0] for pair in pairs] # q is text
+        q_text = batch_q
+        # print(f"q111, {q_text}")
+        
+        
+        # pos
+        # get the score vector
+        batch_p_score_vector = [pair[1] for pair in pairs] # p is score vector
+        batch_p_score_vector = torch.tensor(batch_p_score_vector) # [bs, 4]
+        batch_p_score_vector = batch_p_score_vector[:, :args.vector_dim]
+        # get the text emb
+        batch_p_text_emb = [pair[2] for pair in pairs] # p is text emb
+        batch_p_text_emb = torch.concat(batch_p_text_emb, dim=0) # [bs, 768]
+        # get the symb_enc
+        batch_p_symb_enc = [pair[3] for pair in pairs] # p is symb_enc
+        batch_p_symb_enc = torch.tensor(batch_p_symb_enc) # [bs, 3]
+        
+        
+        # Negative samples
+        batch_n = [random.choices(self.sorted_query2neg[query], k=self.max_neg_candidates) for query in batch_q] # allow duplicates
+        
+        
+        # get the score vector
+        batch_n_score_vector = [pair[0] for sublist in batch_n for pair in sublist]
+        batch_n_score_vector = torch.tensor(batch_n_score_vector) # [bs*100, 4]
+        # reshape to [bs, 100, 4]
+        batch_n_score_vector = batch_n_score_vector.reshape(len(batch_q), self.max_neg_candidates, -1) # [bs, 100, 4]
+        batch_n_score_vector = batch_n_score_vector[:, :, :args.vector_dim]
+        
+        # get the text emb
+        batch_n_text_emb = [pair[1] for sublist in batch_n for pair in sublist] 
+        batch_n_text_emb = torch.concat(batch_n_text_emb, dim=0) # [bs*100, 768]
+        # reshape to [bs, 100, 768]
+        batch_n_text_emb = batch_n_text_emb.reshape(len(batch_q), self.max_neg_candidates, -1) # [bs, 100, 768]
+        
+        # get the symb_enc
+        batch_n_symb_enc = [pair[2] for sublist in batch_n for pair in sublist]
+        batch_n_symb_enc = torch.tensor(batch_n_symb_enc) # [bs*100, 3]
+        # reshape to [bs, 100, 3]
+        batch_n_symb_enc = batch_n_symb_enc.reshape(len(batch_q), self.max_neg_candidates, -1) # [bs, 100, 3]
+            
+        
+        
+        
+        
+        # Create a dictionary for the batch
+        feed_dict = {
+            'query': q_text,
+            'p_score_vector': batch_p_score_vector,
+            'p_text_emb': batch_p_text_emb,
+            'p_symb_enc': batch_p_symb_enc,
+            'n_score_vector': batch_n_score_vector,
+            'n_text_emb': batch_n_text_emb,
+            'n_symb_enc': batch_n_symb_enc,
+            
+        }
+
+
+        return feed_dict
+    
+
+class TestDataset(Dataset):
+    """        
+        data format: {
+            "query": query,
+            "pred_dict": {node_id: score},
+            'score_vector_dict': {node_id: [bm25, bm_25, bm25, ada]},
+            "text_emb_dict": {node_id: text_emb},
+            "ans_ids": [],
+        }
+
+"""
+
+    def __init__(self, saved_data):
+        
+        print(f"Start processing test dataset...")
+        self.text2emb_dict = saved_data['text2emb_dict']
+        self.data = saved_data['data']  
+        
+        self.text_emb_matrix = list(self.text2emb_dict.values())
+        self.text_emb_matrix = torch.concat(self.text_emb_matrix, dim=0)
+        
+        # make the mapping between the key of text2emb_dict and the index of text_emb_matrix
+        self.text2idx_dict = {key: idx for idx, key in enumerate(self.text2emb_dict.keys())}
+        
+        print(f"Complete data preparation: {len(self.data)}")
+        
+             
+        
+      
+    def __len__(self):
+        return len(self.data) 
+
+    def __getitem__(self, idx):
+        
+        # sort the pred_dict by the score
+        pred_dict = self.data[idx]['pred_dict']
+        sorted_ids = sorted(pred_dict.keys(), key=lambda x: pred_dict[x], reverse=True)
+        # get the top 50 candidates
+        sorted_ids = sorted_ids[:50]
+        # get the score vector
+        self.data[idx]['score_vector_dict'] = {key: self.data[idx]['score_vector_dict'][key] for key in sorted_ids}
+        # get the symb_enc_dict
+        self.data[idx]['symb_enc_dict'] = {key: self.data[idx]['symb_enc_dict'][key] for key in sorted_ids}
+        # change from the str to index
+        self.data[idx]['text_emb_dict'] = {key: self.text2idx_dict[value] for key, value in self.data[idx]['text_emb_dict'].items()}
+        self.data[idx]['text_emb_dict'] = {key: self.data[idx]['text_emb_dict'][key] for key in sorted_ids}
+        
+        return self.data[idx]
+    
+    def collate_batch(self, batch):
+        
+        # q
+        batch_q = [batch[i]['query'] for i in range(len(batch))]
+        q_text = batch_q
+        
+        # c
+        batch_c = [list(batch[i]['score_vector_dict'].keys()) for i in range(len(batch))] # [batch, 100]
+        batch_c = torch.tensor(batch_c)
+        # print(f"111, {batch_c.shape}")
+        c_score_vector = [list(batch[i]['score_vector_dict'].values()) for i in range(len(batch))] # [batch, 100, 4]
+        c_score_vector = torch.tensor(c_score_vector)[:, :, :args.vector_dim] # [batch, 100, 4]
+        
+        
+        # print(f"222, {c_vector.shape}")
+        # c_text_emb
+        # c_text_emb = [torch.concat(list(batch[i]['text_emb_dict'].values()), dim=0).unsqueeze(0) for i in range(len(batch))]
+        c_text_emb = [self.text_emb_matrix[list(batch[i]['text_emb_dict'].values())].unsqueeze(0) for i in range(len(batch))]
+        c_text_emb = torch.concat(c_text_emb, dim=0) # [bs, 100, 768]
+        
+        # c_symb_enc
+        c_symb_enc = [list(batch[i]['symb_enc_dict'].values()) for i in range(len(batch))]
+        c_symb_enc = torch.tensor(c_symb_enc) # [bs, 100, 3]
+        
+        
+        # ans_ids
+        ans_ids = [batch[i]['ans_ids'] for i in range(len(batch))] # list of ans_ids
+        
+        # pred_ids
+        pred_ids = batch_c.tolist()
+            
+        
+        # Create a dictionary for the batch
+        feed_dict = {
+            'query': q_text,
+            'c_score_vector': c_score_vector,
+            'c_text_emb': c_text_emb,
+            'c_symb_enc': c_symb_enc,
+            'ans_ids': ans_ids,
+            'pred_ids': pred_ids
+
+        }
+
+        
+        
+        return feed_dict
+    
+
+# ******* loss function ********
+def loss_fn(scores_pos, scores_neg):
+    
+    loss_fct = CrossEntropyLoss(ignore_index=-1)
+
+    # Combine scores
+    scores = torch.cat([scores_pos, scores_neg.squeeze(-1)], dim=1)  # B x (1 + max_neg_candidates*B)
+    # print(f"scores: {scores.shape}")
+    
+    # Create target
+    target = torch.zeros(scores.size(0), dtype=torch.long).to(scores.device)
+
+    # Compute loss
+    loss = loss_fct(scores, target)
+    
+    return loss
+
+
+
+def batch_evaluator(skb, scores_cand, ans_ids, batch):
+
+    results = {}
+    
+    # **** batch wise evaluation ****
+    # evaluate
+    candidates_ids = skb.candidate_ids
+    id_to_idx = {candidate_id: idx for idx, candidate_id in enumerate(candidates_ids)}
+    
+    
+    # initialize the pred_matrix
+    pred_matrix = torch.zeros((scores_cand.shape[0],len(candidates_ids)))
+    
+    
+    # get the index of each pred_ids
+    # flatten the pred_ids
+    flat_pred_ids = torch.tensor(batch['pred_ids']).flatten().tolist()
+    
+    
+    # get the index of each pred_ids
+    pred_idx = [id_to_idx[pred_id] for pred_id in flat_pred_ids]
+    
+    
+    # reshape the pred_idx
+    pred_idx = torch.tensor(pred_idx).reshape(scores_cand.shape[0], -1) # [bs, 100]
+    
+    # move pred_matrix to the device
+    pred_matrix = pred_matrix.to(scores_cand.device)
+    
+    # advanced indexing
+    pred_matrix[torch.arange(scores_cand.shape[0]).unsqueeze(1), pred_idx] = scores_cand.squeeze(-1) # [bs, num_candidates]
+    
+
+    
+    # Create a mapping from candidate IDs to their indices for faster lookup
+    
+
+    # Flatten ans_ids to a single list and map them to indices
+    flat_ans_idx = [id_to_idx[a_id] for sublist in ans_ids for a_id in sublist]
+
+    # Create the row indices for ans_matrix corresponding to the answers
+    row_indices = torch.repeat_interleave(torch.arange(len(ans_ids)), torch.tensor([len(sublist) for sublist in ans_ids]))
+
+    # Create the answer matrix
+    ans_matrix = torch.zeros((scores_cand.shape[0], len(candidates_ids)), device=scores_cand.device)
+    ans_matrix[row_indices, torch.tensor(flat_ans_idx, device=scores_cand.device)] = 1
+
+
+    
+    
+    
+    # batch computing hit1
+    # find the index of the max score
+    max_score, max_idx = torch.max(pred_matrix, dim=1)
+    # check the label of the max idx
+    batch_hit1 = ans_matrix[torch.arange(scores_cand.shape[0]), max_idx]
+    hit1_list = batch_hit1.tolist()
+    
+    
+
+    
+    # batch computing hit@5
+    _, top5_idx = torch.topk(pred_matrix, 5, dim=1)
+    batch_hit5 = ans_matrix[torch.arange(scores_cand.shape[0]).unsqueeze(1), top5_idx]
+    
+    # max with each row
+    batch_hit5 = torch.max(batch_hit5, dim=1)[0]
+    hit5_list = batch_hit5.tolist()
+    
+    
+    
+    # batch computing recall@20
+    _, top20_idx = torch.topk(pred_matrix, 20, dim=1)
+    batch_recall20 = ans_matrix[torch.arange(scores_cand.shape[0]).unsqueeze(1), top20_idx]
+    # sum with each row
+    batch_recall20 = torch.sum(batch_recall20, dim=1)
+    # divide by the sum of the ans_matrix along the row
+    batch_recall20 = batch_recall20 / torch.sum(ans_matrix, dim=1)
+    recall20_list = batch_recall20.tolist()
+    
+    
+    
+    
+    # batch computing mrr
+    # find the highest rank of the answer
+    _, rank_idx = torch.sort(pred_matrix, dim=1, descending=True)
+    # query the answer matrix with the rank_idx
+    batch_mrr = ans_matrix[torch.arange(scores_cand.shape[0]).unsqueeze(1), rank_idx]
+    # find the first rank of the answer
+    batch_mrr = torch.argmax(batch_mrr, dim=1)
+    # add 1 to the rank
+    batch_mrr += 1
+    # divide by the rank
+    batch_mrr = 1 / batch_mrr.float()
+    mrr_list = batch_mrr.tolist()
+    
+    
+    
+
+    results['hit@1'] = hit1_list
+    results['hit@5'] = hit5_list
+    results['recall@20'] = recall20_list
+    results['mrr'] = mrr_list
+    
+
+    
+
+    
+    
+    return results
+    
+
+
+# ***** evaluate *****
+@torch.no_grad()
+def evaluate(reranker, test_loader):
+
+    
+    reranker.eval()
+
+    all_results = {
+        "hit@1": [],
+        "hit@5": [],
+        "recall@20": [],
+        "mrr": []
+    }
+    avg_results = {
+        "hit@1": 0,
+        "hit@5": 0,
+        "recall@20": 0,
+        "mrr": 0
+    }
+    
+    
+    # save the scores and ans_ids, and pred_ids
+    pred_list = []
+    scores_cand_list = []
+    ans_ids_list = []
+    # use tqdm to show the progress
+    for idx, batch in enumerate(tqdm(test_loader, desc='Evaluating', position=0)):
+        batch = move_to_cuda(batch)
+        
+        # Check if the model is wrapped in DataParallel
+        if isinstance(reranker, nn.DataParallel):
+            scores_cand = reranker.module.eval_batch(batch) # q_emb: [bs, 100], c_emb: [bs*100, 100]
+        else:
+            scores_cand = reranker.eval_batch(batch)
+
+
+        # ans_ids
+        ans_ids = batch['ans_ids']
+           
+        results = batch_evaluator(skb, scores_cand, ans_ids, batch)
+                
+        
+        for key in results.keys():
+            all_results[key].extend(results[key])
+            
+        # save the scores and ans_ids, and pred_ids
+        pred_list.extend(batch['pred_ids'])
+        scores_cand_list.extend(scores_cand.cpu().tolist())
+        ans_ids_list.extend(ans_ids)
+        
+    
+    
+    for key in avg_results.keys():
+        avg_results[key] = np.mean(all_results[key])
+    
+    print(f"Results: {avg_results}")
+    
+
+    
+    return avg_results
+    
+    
+# ***** train *****
+def main(train_data, val_data, test_data, skb, dataset_name, args):
+
+
+    epochs = args.epochs
+    device = args.device
+
+    train_size = args.train_batch_size
+    test_size = 64
+
+    train_dataset = TrainDataset(train_data)
+    train_loader = DataLoader(train_dataset, batch_size=train_size, num_workers=32, collate_fn=train_dataset.collate_batch, drop_last=True)
+
+    test_dataset = TestDataset(test_data)
+    test_loader = DataLoader(test_dataset, batch_size=test_size, num_workers=32, collate_fn=test_dataset.collate_batch)
+    
+    val_dataset = TestDataset(val_data)
+    val_loader = DataLoader(val_dataset, batch_size=test_size, num_workers=32, collate_fn=val_dataset.collate_batch)
+    
+
+    # ***** Model *****
+    reranker = PathReranker(socre_vector_input_dim=args.vector_dim, text_emb_input_dim=768, symb_enc_dim=3, args=args)
+    save_dir = f"./data/checkpoints/{dataset_name}/path"
+    os.makedirs(save_dir, exist_ok=True)
+
+    reranker.to(device)
+    # # parallel processing
+    reranker = nn.DataParallel(reranker)
+
+
+    optimizer = torch.optim.Adam(reranker.parameters(), lr=args.lr)
+    best_val_hit1 = float('-inf')
+
+    
+    val_results = evaluate(reranker, val_loader)
+    print(f"Val evaluation")
+    print(val_results)
+    
+    
+    test_results = evaluate(reranker, test_loader)
+    print(f"Test evaluation")
+    print(test_results)
+    
+    # log both val and test results
+    wandb.log({'val_mrr': val_results['mrr'], 'val_hit1': val_results['hit@1'], 'val_hit5': val_results['hit@5'], 'val_recall@20': val_results['recall@20'],
+               'test_mrr': test_results['mrr'], 'test_hit1': test_results['hit@1'], 'test_hit5': test_results['hit@5'], 'test_recall@20': test_results['recall@20']})
+    
+    best_test_results = {}
+    for epoch in tqdm(range(epochs), desc='Training Epochs', position=0):
+        total_loss = 0.0
+        reranker.train()
+        count = 0
+        total_instances = 0
+
+        for batch in tqdm(train_loader):
+            # print(batch)
+            batch = move_to_cuda(batch)
+            # print(batch)
+            
+            scores_pos, scores_neg = reranker(batch)
+            
+            # batch_loss = pairwise_loss(scores_pos, scores_neg)
+            batch_loss = loss_fn(scores_pos, scores_neg)
+            
+            # clear optimizer 
+            optimizer.zero_grad()
+            batch_loss.backward()
+            optimizer.step()
+            
+            # total_loss += batch_loss.item()
+            count += 1 
+            # compute the average loss
+            total_instances += scores_pos.shape[0]
+            total_loss += batch_loss.item()
+                        
+            
+        train_loss = total_loss / total_instances
+        
+        print(f"Epoch {epoch+1}/{epochs}, Average Train Loss: {train_loss}")
+        
+        
+        
+        val_results = evaluate(reranker, val_loader)
+        print(f"Val evaluation")
+        print(val_results)
+        
+        
+        test_results = evaluate(reranker, test_loader)
+        print(f"Test evaluation")
+        print(test_results)
+        
+        # log both val and test results
+        wandb.log({'val_mrr': val_results['mrr'], 'val_hit1': val_results['hit@1'], 'val_hit5': val_results['hit@5'], 'val_recall@20': val_results['recall@20'],
+                'test_mrr': test_results['mrr'], 'test_hit1': test_results['hit@1'], 'test_hit5': test_results['hit@5'], 'test_recall@20': test_results['recall@20'],
+                'train_loss': train_loss})
+        
+        
+        # save the best model when val hit1 is the highest
+        hit1 = val_results['hit@1']
+        if best_val_hit1 < hit1:
+            best_val_hit1 = hit1
+            
+            save_path = f"{save_dir}/best_{best_val_hit1}.pth"
+            
+            if isinstance(reranker, nn.DataParallel):
+                torch.save(reranker.module.state_dict(), save_path)
+            else:
+                torch.save(reranker.state_dict(), save_path)
+            print(f"Checkpoint saved at epoch {epoch+1} with test hits@1 {hit1}")
+            
+            args.checkpoint_path = save_path
+            best_test_results = test_results
+        
+        
+                
+    # save last epoch checkopint
+    save_path = f"{save_dir}/last_{hit1}.pth"
+    if isinstance(reranker, nn.DataParallel):
+        torch.save(reranker.module.state_dict(), save_path)
+    else:
+        torch.save(reranker.state_dict(), save_path)
+    print(f"Final checkpoint saved at {save_path}")
+    
+    
+    
+    # ***** save the results *****
+    results = []
+    results.append(
+        {
+            "config": vars(args),
+            "test_results": best_test_results
+        }
+    )
+    # save the results to json
+    timestamp = time.strftime("%Y%m%d-%H%M%S")
+    output_dir = f"./data/outputs/{dataset_name}"
+    os.makedirs(output_dir, exist_ok=True)
+    with open(f"{output_dir}/results_{timestamp}.json", "w") as f:
+        json.dump(results, f, indent=4)
+        
+    print(best_test_results)
+    
+def get_concat_num(combo):
+    """
+    Determine the value of concat_num based on the combination of embeddings.
+    - score_vec adds +1
+    - text_emb adds +1
+    - symb_enc adds +3
+    """
+    concat_num = 0
+    if combo.get("score_vec", False):  # If score_vec is True
+        concat_num += 1
+    if combo.get("text_emb", False):  # If text_emb is True
+        concat_num += 1
+    if combo.get("symb_enc", False):  # If symb_enc is True
+        concat_num += 3
+        
+        
+    return concat_num
+    
+    
+    
+def parse_args():
+    
+    parser = argparse.ArgumentParser(description="Run Pathreranker with dynamic combinations of embeddings.")
+    
+    # Add arguments for model configurations
+    parser.add_argument("--train_batch_size", type=int, default=256, help="Batch size for training or evaluation.")
+    parser.add_argument("--lr", type=float, default=3e-5, help="Learning rate for optimizer.")
+    parser.add_argument("--epochs", type=int, default=100, help="Number of epochs to train the model.")
+    parser.add_argument("--device", type=str, default="cuda", help="Device to run the model (e.g., 'cuda' or 'cpu').")
+    
+    # Add arguments for the dataset
+    parser.add_argument("--dataset_name", type=str, default="prime", help="Name of the dataset to use.")
+    # paths
+    parser.add_argument("--train_path", type=str, default=f"../prime_train.pkl", help="Path to the training data.")
+    parser.add_argument("--test_path", type=str, default=f"../prime_test.pkl", help="Path to the test data.")
+    parser.add_argument("--val_path", type=str, default=f"../prime_val.pkl", help="Path to the validation data.")
+    
+    # add concat_num
+    parser.add_argument("--concat_num", type=int, default=0, help="Number of concatenation of embeddings.")
+    
+    # checkpoint save path
+    parser.add_argument("--checkpoint_path", type=str, default="", help="Path to save the checkpoints.")
+    parser.add_argument("--vector_dim", type=int, default=4, help="Dimension of the similarity vector.")
+    
+    
+    # Parse the base arguments
+    args = parser.parse_args()
+    return args
+    
+    
+if __name__ == "__main__":
+    
+    
+    base_args = parse_args()
+    
+    dataset_name = base_args.dataset_name
+    train_path = base_args.train_path
+    test_path = base_args.test_path
+    val_path = base_args.val_path
+    
+    
+    with open(test_path, "rb") as f:
+        test_data = pkl.load(f)
+        
+    with open(train_path, "rb") as f:
+        train_data = pkl.load(f)
+        
+    with open(val_path, "rb") as f:
+        val_data = pkl.load(f)
+    
+    # load skb
+    skb = load_skb(dataset_name)
+    
+    # set all
+    combo = {
+        "text_emb": True,
+        "score_vec": True,
+        "symb_enc": True
+    }
+    concat_num = get_concat_num(combo)
+    
+    wandb.init(project=f'Reranking-{dataset_name}', name=f"path")
+    args = argparse.Namespace(**vars(base_args), **combo)
+    args.concat_num = concat_num
+    
+    main(train_data, val_data, test_data, skb, dataset_name, args)
+        

Reranking/utils.py

@@ -0,0 +1,90 @@
+import sys
+from pathlib import Path
+# Get the absolute path of the current script
+current_file = Path(__file__).resolve()
+project_root = current_file.parents[1]
+# Add the project root to the system path
+sys.path.append(str(project_root))
+
+import random
+import torch
+import os
+from stark_qa.evaluator import Evaluator
+import torch.nn as nn
+from typing import Any, Union, List, Dict
+
+
+class ModelForSTaRKQA(nn.Module):
+    
+    def __init__(self, skb, query_emb_dir='.'):
+        """
+        Initializes the model with the given knowledge base.
+        
+        Args:
+            skb: Knowledge base containing candidate information.
+        """
+        super(ModelForSTaRKQA, self).__init__()
+        self.skb = skb
+
+        self.candidate_ids = skb.candidate_ids
+        self.evaluator = Evaluator(self.candidate_ids)
+        
+    def evaluate(self, 
+                 pred_dict: Dict[int, float], 
+                 answer_ids: Union[torch.LongTensor, List[Any]], 
+                 metrics: List[str] = ['mrr', 'hit@3', 'recall@20'], 
+                 **kwargs: Any) -> Dict[str, float]:
+        """
+        Evaluates the predictions using the specified metrics.
+        
+        Args:
+            pred_dict (Dict[int, float]): Predicted answer ids or scores.
+            answer_ids (torch.LongTensor): Ground truth answer ids.
+            metrics (List[str]): A list of metrics to be evaluated, including 'mrr', 'hit@k', 'recall@k', 
+                                 'precision@k', 'map@k', 'ndcg@k'.
+                             
+        Returns:
+            Dict[str, float]: A dictionary of evaluation metrics.
+        """
+        return self.evaluator(pred_dict, answer_ids, metrics)
+    
+    def evaluate_batch(self, 
+                pred_ids: List[int],
+                pred: torch.Tensor, 
+                answer_ids: Union[torch.LongTensor, List[Any]], 
+                metrics: List[str] = ['mrr', 'hit@3', 'recall@20'], 
+                **kwargs: Any) -> Dict[str, float]:
+        return self.evaluator.evaluate_batch(pred_ids, pred, answer_ids, metrics)
+    
+
+def seed_everything(seed=0):
+    random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+    os.environ['PYTHONHASHSEED'] = str(seed)
+    torch.backends.cudnn.deterministic = True
+    torch.backends.cudnn.benchmark = False
+    
+
+def move_to_cuda(sample):
+    if len(sample) == 0:
+        return {}
+
+    def _move_to_cuda(maybe_tensor):
+        if torch.is_tensor(maybe_tensor):
+            return maybe_tensor.cuda()
+        elif isinstance(maybe_tensor, dict):
+            return {
+                key: _move_to_cuda(value)
+                for key, value in maybe_tensor.items()
+            }
+        # elif isinstance(maybe_tensor, list):
+        #     return [_move_to_cuda(x) for x in maybe_tensor]
+        else:
+            return maybe_tensor
+
+    return _move_to_cuda(sample)
+
+
+if __name__ == "__main__":
+    print("Testing Utils")