codetext.txt

### Start of file: utils.py ###
import os
import requests
from concurrent.futures import ThreadPoolExecutor
import time
import queue
import concurrent
import threading
import random
import copy
from collections import deque
import json
from transformers import AutoModelForCausalLM, AutoTokenizer
from transformers import pipeline
import torch
import gc
from unidecode import unidecode
import re
import math

FEW_SHOT_EXAMPLES = """
Example questions and correct answers:
Common Context: entity_B is the son of entity_A. entity_E is the sister of entity_A. entity_B leads entity_C. Entity_D is a member of Entity_C. Entity_D is a friend of entity_E. entity_E has mother entity_F who likes the services of entity_C.
question 1: entity_A->(father of)->(leader of)->?
Options: 1. entity_F,\n 2. entity_C,\n 3. entity_D,\n 4. entity_E,\n 5. entity_B
answer: 2. entity_C
explanation: entity_A->(father of)entity_B->(leader of)entity_C
how to get answer: find who entity_A is father of to get entity_B, then find what B is the leader of to get entity_C which the final answer.

question 2: entity_B->(chief of)->(constitues)->(companion of)->?
Options: 1. entity_F,\n 2. entity_C,\n 3. entity_D,\n 4. entity_E,\n 5. entity_A
answer: 4. entity_E
explanation: entity_B->(chief of)entity_C->(constitues)entity_D->(companion of)entity_E
how to get answer: find what entity_B is the chief of to get entity_C, find what entity_C constitutes of to get entity_D, then find the companion of entity_D of to get entity_E.
"""

# question 3: entity_C->(has membership of)->(friend)->(mom)->?
# Options: 1. entity_F,\n 2. entity_C,\n 3. entity_D,\n 4. entity_E,\n 5. entity_A
# answer: 1. entity_F
# explanation: entity_C->(has membership of)entity_D->(friend)entity_E->(mom)entity_F
# how to get answer: find who entity_C has membership of to get entity_D, find who is the friend of entity_D to get entity_E, then find the mom of entity_E to get entity_F.

# question 4: entity_A->(sister)->(mother)->?
# Options: 1. entity_F,\n 2. entity_C,\n 3. entity_D,\n 4. entity_E,\n 5. entity_A
# answer: 5. entity_F
# explanation: entity_A->(sister)entity_E->(mother)entity_F
# how to get answer: find who entity_A's sister is to get entity_E, then find the mother of entity_E to get entity_F.

# question 5: entity_A->(son)->(leader)->(member)->(friend)?
# Options: 1. entity_F,\n 2. entity_C,\n 3. entity_D,\n 4. entity_E,\n 5. entity_A
# answer: 3. entity_E
# explanation: entity_A->(son)entity_B->(leader)entity_C->(member)entity_D->(friend)entity_E
# how to get answer: find who entity_A's son is to get entity_B, then find the leader of entity_B to get entity_C, find the member of entity_C to get entity_D, then find the friend of entity_D to get entity_E.

LLM_PROMPT_TEMPLATE = """
{few_shot_examples}

Actual Query:
Given Context:
{context}

Answer the question:
{query}

answer the question by selecting the correct answer from the following options:
{options}

The format for beginning your response is:
correct answer: <option_number>. <answer>, because <succinct reason>

follow this exact format and only choose from the given options
"""

# CHECKER_INITIAL_PROMPT = f"""
# You are a correct answer evaluator. Your inputs will consist of a question and a correct answer, and a answer from a model.
# The questions will be of the form:
# {FEW_SHOT_EXAMPLES}

# Your response should start with 'correct' or 'wrong' based on whether the model's answer means the correct answer in both technical and semantic terms.

# Start response with 'correct' or 'wrong' only and nothing else. Then explain the reasons.
# """

def sort_vertices_by_outdegree(graph):
    """
    Sorts the vertices in the graph based on their outdegree in descending order.
    
    :param graph: A dict of dicts representing the graph, where graph[node_start] = {node1: edge_weight, node2: edge_weight, ...}
    :return: A list of vertices sorted by their outdegree in descending order.
    """
    # Calculate outdegree for each vertex
    outdegree = {vertex: len(edges) for vertex, edges in graph.items()}
    
    # Sort vertices based on outdegree in descending order
    sorted_vertices = sorted(outdegree, key=outdegree.get, reverse=True)
    
    return sorted_vertices

def sort_vertices_by_measure(graph, k, weights):
    """
    Sorts the vertices in the graph based on a custom measure. 

    The measure considers the number of neighbors at different distances (up to k) and weights assigned to each distance.

    :param graph: A dict of dicts representing the graph, where graph[node_start] = {node1: edge_weight, node2: edge_weight, ...}
    :param k: The maximum distance to consider for neighbors.
    :param weights: A dictionary mapping distances to weights.
    :return: A list of vertices sorted by the calculated measure in descending order.
    """
    # Calculate the measure for each vertex
    measure = {vertex : 0 for vertex in graph.keys()}
    num_neighbors = {(vertex, 0): 1 for vertex in graph.keys()}
    for i in range(1, k+1):
        for vertex, ents_rels in graph.items():
            key = (vertex, i)
            total = 0
            for neighbor in ents_rels.keys():
                total += num_neighbors.get((neighbor, i-1), 1)
            num_neighbors[key] = total
    
    for key, value in num_neighbors.items():
        if key[1] <= 0:
            continue
        if num_neighbors[(key[0], key[1]-1)] == 0:
            continue
        measure[key[0]] += (value/num_neighbors[(key[0], key[1]-1)]) * weights[key[1]]
    # Sort vertices based on the measure
    sorted_vertices = sorted(measure, key=measure.get, reverse=True)
    
    return sorted_vertices

def generate_answer_options(correct_ans, distractors, path_entities, random_entities, wikidata_util,
                            min_num_options=4):
    # assumption: distractors are ordered by decreasing distracting power
    # assumption: path_entities is also ordered by increasing proximity to the correct answer
    #random entities: list[list[str]], 0th element: list[ents] that are from path[-1], so on
    #   A
    #  / \r
    # B   E
    # |r
    # C
    assert correct_ans == path_entities[0]
    options = [(correct_ans, path_entities[1])]
    options.extend(distractors)
    if len(options) >= min_num_options:
        options = options[:min_num_options]
        random.shuffle(options)
        return options
    if len(path_entities) > 0:
        options.extend([(path_entities[i], path_entities[i+1]) for i in range(1, len(path_entities)-1)])
    if len(options) >= min_num_options:
        options = options[:min_num_options]
        random.shuffle(options)
        return options
    rel_to_find = wikidata_util[path_entities[1]][path_entities[0]]

    good_random_entities = [] # has the same relation as that connecting correct answer to its parent

    for i in range(1, len(random_entities)):
        parent_ent = path_entities[i]
        for pos_ent in random_entities[i]:
            if i > 0:
                if pos_ent == path_entities[i-1]:
                    continue
            if wikidata_util[parent_ent][pos_ent] == rel_to_find:
                good_random_entities.append((pos_ent, parent_ent))
    options.extend(good_random_entities)
    num_random = 2
    random_options = []
    for i in range(1, len(random_entities)):
        parent_ent = path_entities[i]
        pos_ents = [(ent, parent_ent) for ent in random_entities[i] if (ent, parent_ent) not in options]
        random_options.extend(random.sample(pos_ents, min(num_random, len(pos_ents))))  
    options.extend(random_options)
    if len(options) >= min_num_options:
        options = options[:min_num_options] #ensure correct answer is in the options as it is at index 0
        random.shuffle(options)
    else:
        #in case we still don't have enough options, override num_random on origin vertex to get more options
        origin_vertex = path_entities[-1]
        pos_ents = [(ent, origin_vertex) for ent in random_entities[-1] if (ent, origin_vertex) not in options]
        options_needed = min(min_num_options - len(options), len(pos_ents)) #ensure we only have min_num_options options
        options.extend(random.sample(pos_ents, options_needed))
        random.shuffle(options) 
    return options
        

        

class GraphAlgos():
    def __init__(self, graph: dict, entity_aliases:dict, relation_aliases:dict) -> None:
        """
        Initializes the GraphAlgos object with the given graph and alias dictionaries.

        :param graph: A dict of dicts representing the graph.
        :param entity_aliases: A dictionary mapping entity IDs to lists of aliases.
        :param relation_aliases: A dictionary mapping relation IDs to lists of aliases.
        """
        all_vertices = set(graph.keys()) | {neighbor for neighbors in graph.values() for neighbor in neighbors.keys()}
        self.graph = {vertex: graph.get(vertex, {}) for vertex in all_vertices}
        self.rel2ent_graph = {vertex: {rel:[]  for rel in neighbors.values()} for vertex, neighbors in graph.items()}
        for vertex, neighbors in self.graph.items():
            if vertex not in self.rel2ent_graph:
                assert self.graph[vertex] == {}, f"Vertex {vertex} in graph but not in rel2ent_graph"
                self.rel2ent_graph[vertex] = {}
                continue
            for neighbor, rel in neighbors.items():
                self.rel2ent_graph[vertex][rel].append(neighbor)
        self.entity_aliases = entity_aliases
        self.relation_aliases = relation_aliases
    def bfs(self, start):
        """
        Performs Breadth-First Search (BFS) on the graph starting from the specified vertex.

        :param start: The starting vertex for the BFS.
        :return: A dictionary mapping vertices to their distances from the starting vertex.
        """
        distances = {vertex: -float('infinity') for vertex in self.graph}
        distances[start] = 0
        queue = deque([start])

        while queue:
            current_vertex = queue.popleft()
            for neighbor in self.graph[current_vertex].keys():
                if distances[neighbor] == -float('infinity'):
                    distances[neighbor] = distances[current_vertex] + 1
                    queue.append(neighbor)

        return distances

    def compute_diameter(self):
        max_distance = 0
        for vertex in self.graph:
            distances = self.bfs(vertex)
            max_distance = max(max_distance, max(distances.values()))

        return max_distance

    def dfs(self, start, visited=None):
        if visited is None:
            visited = set()
        visited.add(start)

        for next in self.graph[start].keys():
            if next not in visited:
                continue
            self.dfs(next, visited)
        return visited
    
    def dfs_path(self, start, length, path=None):
        if path is None:
            path = [start]

        if len(path) - 1 == length:  # -1 because the length of path includes the starting node
            return path  # Return the path if it's of the desired length
        if start not in self.graph:
            return None
        neighbors = self.graph[start]
        if not neighbors:
            return None # Return None if there are no neighbors
        neighbors_sample = random.sample(neighbors.keys(), len(neighbors))
        for neighbor in neighbors_sample:
            if neighbor not in path:
                extended_path = path + [neighbor]
                result = self.dfs_path(neighbor, length, extended_path)
                if result:
                    return result  # Return the first path of the desired length found
        return None  # Return None if no path of the desired length is found
    
    def get_vertices(self):
        return list(self.graph.keys())

    def get_relations(self):
        return list({relation for neighbors in self.graph.values() for relation in neighbors.values()})

    def get_relation_for_vertex(self, start_vertex, target_vertex):
        if target_vertex in self.graph[start_vertex]:
            return self.graph[start_vertex][target_vertex]
        return None
    
    def get_path_for_vertices(self, start, end, k=5):
        """
        Finds a path between two vertices in the graph using Breadth-First Search (BFS) with a limited depth.

        :param start: The starting vertex for the path search.
        :param end: The target ending vertex for the path.
        :param k: The maximum depth to explore in the BFS (default: 5).
        :return: A list of vertices representing the found path, or None if no path is found within the specified depth.
        """
        queue = deque([[start]])
        visited = set()
        while queue:
            path = queue.popleft()
            vertex = path[-1]
            if vertex == end:
                return path
            elif vertex not in visited and len(path) < k:
                for neighbor in self.graph[vertex].keys():
                    new_path = list(path)
                    new_path.append(neighbor)
                    queue.append(new_path)
                visited.add(vertex)
        return None
    
    def get_queries_for_relpath(self, rel_path: list, start_vertex: str) -> list:
        """
        Finds all possible paths in the graph that follow the given sequence of relations, starting from the specified vertex. 

        This function recursively explores the graph, considering neighbors connected by the relations in the `rel_path` list.

        :param rel_path: A list of relations representing the desired path.
        :param start_vertex: The starting vertex for the path search.
        :return: A list of lists, where each inner list represents a valid path found in the graph following the specified relations.
        """
        if len(rel_path) == 0:
            return [[start_vertex]]
        rel_use = rel_path[0]
        neighbors = self.rel2ent_graph[start_vertex].get(rel_use, [])
        if len(neighbors) == 0:
            return []
        possible_paths = []
        for neighbor in neighbors:
            possible_neighbor_paths = self.get_queries_for_relpath(rel_path[1:], neighbor)
            if len(possible_neighbor_paths) == 0:
                continue
            for path in possible_neighbor_paths:
                possible_paths.append([start_vertex] + path)
        return possible_paths
    
    def generate_query_for_path(self, path):
        def query_path_aux(path):
            if len(path) == 1:
                entity_alias = get_alias(path[0], self.entity_aliases)
                return entity_alias, entity_alias
            
            rel = self.get_relation_for_vertex(path[-2], path[-1])
            rel_alias = get_alias(rel, self.relation_aliases)
            rest_query, entity_alias = query_path_aux(path[:-1])
            query = f"{rest_query}->({rel_alias}) "
            return query, entity_alias
            
        query = ""
        rest_query, entity_alias = query_path_aux(path)
        query += rest_query + '->?'
        return query, entity_alias, len(path)
    
    def generate_query_for_vertices(self, start, end, k=5, path=None):
        if path is None:
            path = self.get_path_for_vertices(start, end, k)
        if path is None:
            return None
        return self.generate_query_for_path(path)
    
    def sample_random_vertex(self, vertex_list=None):
        if vertex_list is None:
            vertex_list = list(self.graph.keys())
        return random.choice(vertex_list)
    
    def generate_random_path(self, path_len=25, source=None):
        """
        Generates a random path in the graph with the specified length.

        The generated path attempts to be unique, meaning there should be only one possible path 
        corresponding to the sequence of relations in the path.

        :param path_len: The desired length of the path.
        :param source: (Optional) The starting vertex for the path. If not provided, a random vertex is chosen.
        :return: A list of vertices representing the generated path, or None if a unique path cannot be found within 500 attempts. 
        """
        path = None
        i = 0
        while path is None and i < 500:
            if source is None:
                start = self.sample_random_vertex()
            else:
                start = source
            path = self.dfs_path(start, path_len)
            #check if path in unique
            if path is not None:
                rel_path = [self.get_relation_for_vertex(path[i], path[i+1]) for i in range(len(path)-1)]
                # if len(self.rel2ent_graph[path[-2]][rel_path[-1]]) > 1:
                #     path = None #answer should be unique
                #     continue
                possible_paths = self.get_queries_for_relpath(rel_path, start)
                assert len(possible_paths) > 0, f"No possible paths found for {start} -> {path[-1]}"
                if len(possible_paths) > 1:
                    path = None #not unique
            i += 1
        return path
    
    def get_best_distractor(self, start_vertex, path, do_choose=True):
        """
        Gets the best distractor vertex for a given path.
        Assumes that the path is unique. If the path is not unique, the behavior is undefined.

        Returns a Union(tuple(string, string), list(tuple(string, string)))
        the first string is the distractor and the second string in the node in the path that leads to the distractor
        do_choose=True returns a single tuple, do_choose=False returns a list of tuples of possible distractors
        """
        rel_path = [self.get_relation_for_vertex(path[i], path[i+1]) for i in range(len(path)-1)]
        best_distractor = None
        all_distractors_pos = []
        distractor_weights = []
        for i in range(len(rel_path)-2, -1, -1):
            cur_vertex = path[i]
            neighbors = self.rel2ent_graph[cur_vertex].get(rel_path[i], [])
            if len(neighbors) == 0:
                continue
            pos_distractors = [neighbor for neighbor in neighbors if neighbor not in path]
            for distractor in pos_distractors:
                all_distractors_pos.append((distractor, cur_vertex))
                distractor_weights.append(i)
                
        if not do_choose:
            return all_distractors_pos
        if len(all_distractors_pos) > 0:
            best_distractor_index = random.choices(range(len(all_distractors_pos)), weights=distractor_weights, k=1)[0]
            # print(f"Best Distractor Weight: {distractor_weights[best_distractor_index]}")
            best_distractor = all_distractors_pos[best_distractor_index]
        if best_distractor is not None:
            assert best_distractor not in path, f"Best distractor {best_distractor} is in path {path}"
        return best_distractor
    
    
    def generate_random_query(self, k=5, return_path=False, source=None):
        """
        Generates a random query by sampling a random path in the graph.

        :param k: The maximum path length.
        :param return_path: If True, returns the path used to generate the query.
        :param source: Optional starting vertex for the path.
        :return: A tuple containing the generated query, starting vertex, correct answer(s), and optionally the path.
        """
        path_len = random.randint(1, k)
        path = self.generate_random_path(path_len, source)
        if path is None:
            return None
        # print(f"Query Path: {str(path[start_idx:path.index(end)+1])}")
        start = path[0]
        end = path[-1]
        rel_correct = self.get_relation_for_vertex(path[-2], path[-1])
        # print(rel_correct, "rel correct")
        correct_answers = self.rel2ent_graph[path[-2]][rel_correct]
        if return_path:
            return self.generate_query_for_vertices(start, end, k, path), start, correct_answers, path
        return self.generate_query_for_vertices(start, end, k, path), start, correct_answers, None
    
    def create_subgraph_within_radius(self, start_vertex, k):
        """
        Creates a subgraph containing all vertices within a radius of k from the given start_vertex.
        :param start_vertex: The vertex from which the radius is measured.
        :param k: The radius within which vertices will be included in the subgraph.
        :return: A subgraph as a dict of dicts including all vertices within the specified radius.
        """
        visited = {start_vertex: 0}  # Tracks visited vertices and their levels
        queue = [(start_vertex, 0)]  # Queue for BFS, storing tuples of (vertex, level)
        subgraph = {}  # The resulting subgraph
        
        while queue:
            current_vertex, level = queue.pop(0)
            
            if level > k:
                break  # Stop if the current level exceeds k
            
            # Initialize subgraph entry for the current vertex if not already present
            if current_vertex not in subgraph:
                subgraph[current_vertex] = {}
            
            # Explore neighbors
            for neighbor, weight in self.graph.get(current_vertex, {}).items():
                if neighbor not in visited or visited[neighbor] > level + 1:
                    visited[neighbor] = level + 1
                    queue.append((neighbor, level + 1))
                    
                    # Add the neighbor to the subgraph with the corresponding edge weight
                    if level + 1 <= k:
                        subgraph[current_vertex][neighbor] = weight
                        
        return subgraph
    def get_best_vertices(self, num=1000, method='outdegree', **kwargs):
        """
        Gets the best vertices from the graph based on the specified method.
        :param num: The number of vertices to return.
        :param method: The method to use for selecting the best vertices. Supported methods are:
            - 'outdegree': Selects vertices with the highest outdegree.
            - 'measure': Selects vertices based on a custom measure considering neighbor distances and weights.
        :param kwargs: Additional keyword arguments for the chosen method. 
            For 'measure':
                - k (int, optional): The maximum distance to consider for neighbors (default: 4).
                - weights (dict, optional): A dictionary mapping distances to weights (default: {1: 0.75, 2: 1.25, 3: 1, 4: 0.25}).
        :return: A list of the best vertices, sorted according to the chosen method.
        :raises ValueError: If an unsupported method is specified.
        """
        graph_copy = self.graph.copy()
        if method == 'outdegree':
            vertices_out = sort_vertices_by_outdegree(graph_copy)
        elif method == 'measure':
            k = kwargs.get('k', 4)
            weights = kwargs.get('weights', {1: 0.75, 2: 1.25, 3: 1, 4: 0.25})
            vertices_out = sort_vertices_by_measure(graph_copy, k, weights)
        else:
            raise ValueError(f"Method {method} not supported")
        return vertices_out[:num]

class MistralChecker():
    def __init__(self, model_path, device='cuda:1') -> None:
        """
        Checks the correctness of a model answer using the Mistral language model.

        :param prompt: The prompt containing the question, ground truth answer, and model answer.
        :param return_response: If True, returns the raw response from Mistral.
        :return: A tuple containing 1 if the answer is correct, 0 otherwise, and optionally the raw response from Mistral.
        """
        self.device = device
        self.model = AutoModelForCausalLM.from_pretrained(model_path, torch_dtype=torch.float16).to(self.device)
        self.tokenizer = AutoTokenizer.from_pretrained(model_path)
        self.generation_config = {'temperature':0}
        self.count = 0
    def form_prompt(self, question, correct_ans, model_ans):
        return f"Question: {question} Ground Truth: {correct_ans}. Model Answer: {model_ans}"

    def checker(self, prompt: str, return_response=False):
        """
        input in form: "Question: What is David Beckham's National
        ity? Ground Truth: England. Model Answer: British"
        returns 1 for yes 0 for no and response if return response = True
        """
        
        messages = [
            {"role": "user", "content": f"{CHECKER_INITIAL_PROMPT}"},
            {"role": "assistant", "content": f"Okay, I will evaluate the correctness of the answers and the reasoning of the model_answer with respect to the correct answers technically and semantically."},
            {"role": "user", "content": f"Evaluate Input data: [{prompt}] "}
        ]
        encodeds = self.tokenizer.apply_chat_template(messages, return_tensors="pt")

        model_inputs = encodeds.to(self.device)
        generated_ids = self.model.generate(model_inputs, max_new_tokens=2, do_sample=False, pad_token_id=self.tokenizer.eos_token_id, temperature=0.0)
        decoded = self.tokenizer.batch_decode(generated_ids[:, model_inputs.size(1):])
        print(decoded)
        decoded = str(decoded[0]).strip().lower().strip()
        correct_ans = int('correct' in decoded)
        del model_inputs, generated_ids
        if return_response:
            return correct_ans, decoded
        return correct_ans, None

    def raw_checker(self, question, correct_ans, model_ans, return_response=False):
        prompt = self.form_prompt(question, correct_ans, model_ans)
        return self.checker(prompt, return_response)

def get_alias(id, aliases):
    """
    Retrieves a random alias for the given ID from the provided aliases dictionary.

    :param id: The ID for which to retrieve an alias.
    :param aliases: A dictionary mapping IDs to lists of aliases.
    :return: A randomly chosen alias from the list associated with the ID.
    """
    aliases_choices = aliases[id]
    return random.choice(aliases_choices)
    
def form_alias_question(question, path, entity_aliases, relation_aliases, entity_name2id, relation_name2id, graph_algos):
    """
    Transforms a question by replacing entities and relations with their aliases.

    Redundant, Not being used because of errors in name2id
    """
    entity_replace = path[0]
    entity_id = entity_name2id[entity_replace]
    entity_alias = get_alias(entity_id, entity_aliases)
    entity_alias = unidecode(entity_alias)
    question = question.replace(entity_replace, entity_alias)
    for i in range(len(path)-1):
        relation_name = graph_algos.get_relation_for_vertex(path[i], path[i+1])
        relation_name = unidecode(relation_name).lower()
        try:
            relation_id = relation_name2id[relation_name]
            try:
                relation_alias = get_alias(relation_id, relation_aliases)
            except KeyError:
                try:
                    relation_alias = get_alias(relation_id, entity_aliases) #relations maybe entities too(eg. country)
                except KeyError:
                    relation_alias = relation_name
                    print(f"Relation {relation_name, relation_id} not found in aliases using default")
        except KeyError:
            relation_alias = relation_name
            print(f"Relation {relation_name} not found in name2id using default")
        relation_alias = unidecode(relation_alias)
        question = question.replace(relation_name, relation_alias)

    return question, entity_alias

def load_aliases(path):
    """
    Loads aliases from a file.

    :param path: The path to the file containing aliases.
    :return: A dictionary mapping entity/relation IDs to lists of aliases. 
    """
    possible_entities = {}
    with open(path, 'r') as f:
        for line in f:
            line = line.strip()
            line = line.split('\t')
            line = [x.strip() for x in line]
            possible_entities[line[0]] = line[1:]
    return possible_entities

def form_context_list(query_path, wikidata_text_edge, wikidata_util, entity_top_alias):
    context_list = []
    for i in range(len(query_path)-1):
        source, dest = query_path[i], query_path[i+1]
        if dest not in wikidata_text_edge[source]:
            print(source, dest)
            return None
        relevant_text = wikidata_text_edge[source][dest]
        if wikidata_util[source][dest] == 'P20': # death place
            relevant_text = [f"{entity_top_alias[source]} died at {entity_top_alias[dest]}"]
        else:
            relevant_text = []
        relevant_text.extend(wikidata_text_edge[source][dest])
        # if type(relevant_text) == list:
        #     relevant_text = ' '.join(relevant_text)
        context_list.append(relevant_text)
    return context_list

# def simple_checker(model_answer, correct_answer, correct_answer_aliases, correct_id):
#     """
#     Performs a simple check to see if the model answer is correct.

#     Checks if the correct answer or any of its aliases are present in the model answer (case-insensitive).

#     :param model_answer: The answer generated by the model.
#     :param correct_answer: The ground truth correct answer.
#     :param correct_answer_aliases: A dictionary mapping correct answer IDs to their aliases.
#     :param correct_id: The ID of the correct answer.
#     :return: 1 if the model answer is considered correct, 0 otherwise.
#     """
#     model_answer = unidecode(model_answer).lower()
#     correct_answer = unidecode(correct_answer).lower()
#     pattern = r'\b' + re.escape(correct_answer) + r'\b'
#     matches = re.search(pattern, model_answer)
#     if matches:
#         return 1

#     if correct_id not in correct_answer_aliases:
#         return 0
#     for answer_alias in correct_answer_aliases[correct_id]:
#         pattern = r'\b' + re.escape(unidecode(answer_alias).lower()) + r'\b'
#         matches = re.search(pattern, model_answer)
#         if matches:
#             return 1
#     return 0

def dumb_checker(model_answer, correct_answer_num):
    """
    Performs a simple check to see if the model answer is correct.

    Checks if the correct answer or any of its aliases are present in the model answer (case-insensitive).

    :param model_answer: The answer generated by the model.
    :param correct_answer_num: The ground truth correct answer number.
    :return: 1 if the model answer is considered correct, 0 otherwise.
    """
    model_answer = unidecode(model_answer).lower()
    correct_answer_num = unidecode(str(correct_answer_num)).lower()
    pattern = r'\bcorrect answer: \s*\**\s*[\[<{(]?' + re.escape(correct_answer_num) + '\s*\**\s*[\]>})]?[.]?'
    matches = re.search(pattern, model_answer)
    if matches:
        return 1
    return 0

def create_context_list(all_sents, relevant_sents_path, relevant_sents_opts, tokenizer, max_length=15000):
    # Flatten relevant_sents and calculate its total tokenized length
    if tokenizer is None:
        #for API only models
        combined_list = []
        all_sents_set = set()
        for a_sublist in all_sents:
            add_list = []
            for a_item in a_sublist:
                if a_item not in add_list and a_item not in all_sents_set:
                    add_list.append(a_item)
                    all_sents_set.add(a_item)
            if len(add_list) > 0:
                combined_list.append(add_list)
        return combined_list
    flat_relevant_sents = [item for sublist in relevant_sents_path for item in sublist]
    tokenized_relevant_sents = tokenizer(flat_relevant_sents, add_special_tokens=False)
    total_length_relevant_sents = sum(len(ids) for ids in tokenized_relevant_sents['input_ids'])

    if total_length_relevant_sents > max_length:
        print(f"Total length of relevant sentences ({total_length_relevant_sents}) exceeds maximum length ({max_length}).")
        return []

    option_sentences_add = []
    # Tokenize option sentences if we can
    for option_sents in relevant_sents_opts:
        tokenized_relevant_sents = tokenizer(option_sents, add_special_tokens=False)
        len_tokens = sum(len(ids) for ids in tokenized_relevant_sents['input_ids'])
        if len_tokens + total_length_relevant_sents <= max_length:
            option_sentences_add.append(option_sents)
    
    relevant_sents = relevant_sents_path + option_sentences_add
    flat_relevant_sents = [item for sublist in relevant_sents for item in sublist]
    set_relevant_sents = set(flat_relevant_sents)
    to_include = set()
    current_length = total_length_relevant_sents
    seen_strings = set(flat_relevant_sents)  # To track duplicates and included strings

    # Check each item in all_sents for inclusion
    for a_sublist in all_sents:
        for a_item in a_sublist:
            item_tokenized_length = len(tokenizer(a_item, add_special_tokens=False))
            if a_item in set_relevant_sents or (a_item not in seen_strings and current_length + item_tokenized_length <= max_length):
                to_include.add(a_item)
                seen_strings.add(a_item)
                if a_item not in set_relevant_sents:  # Only add to current_length if not already accounted for
                    current_length += item_tokenized_length

    combined_list = []
    all_sents_set = set()
    # Ensure each item is only added once
    for a_sublist in all_sents:
        add_list = []
        for a_item in a_sublist:
            if a_item in to_include and a_item not in add_list and a_item not in all_sents_set:
               add_list.append(a_item)
               all_sents_set.add(a_item)
        if len(add_list) > 0:
            combined_list.append(add_list)
    return combined_list

def get_all_context(query_path, wikidata_text_sentencized):
    all_context = []
    for entity in query_path:
        if entity in wikidata_text_sentencized:
            all_context.append(wikidata_text_sentencized[entity])
    return all_context
        
def get_random_entities(query_path, wikidata_util):
    random_entities = []
    for i in range(len(query_path)):
        parent_ent = query_path[i]
        pos_ents = list(wikidata_util[parent_ent].keys())
        random_entities.append(pos_ents)
    return random_entities


def get_query_data(graph_algos, source, id2name, graph_text_edge, graph_text_sentencized, tokenizer, distractor_query=False, k=5, shuffle_context=True, max_context_length=30000):
    while True:
        distractor=None
        node_distracted=None
        query_results = graph_algos.generate_random_query(k, return_path=True, source=source) # allow sampling with replacement
        if query_results is None:
            return None
        all_distractors = []
        if distractor_query:
            distractor_tuple = graph_algos.get_best_distractor(query_results[1], query_results[3])
            if distractor_tuple is None:
                continue
            distractor, node_distracted = distractor_tuple
            all_distractors = graph_algos.get_best_distractor(query_results[1], query_results[3], do_choose=False)
            
        query_inf, _, correct_ids, ids_path = query_results
        query, entity_alias, k_num = query_inf
        path = [id2name[ids_path[i]] for i in range(len(ids_path))]
        if 'country of' in query:
            query = query.replace('country of', 'country location of') # to avoid ambiguity
        true_ids_path = ids_path.copy()
        # if not distractor_query:
        #     random_distractor_parent = random.choice(list(graph_text_edge.keys()))
        #     try:
        #         random_distractor = random.choice(list(graph_text_edge[random_distractor_parent].keys()))
        #     except:
        #         random_distractor = None
        #     if random_distractor not in true_ids_path and random_distractor is not None:
        #         distractor = random_distractor
        #         node_distracted = random_distractor_parent
        options = generate_answer_options(true_ids_path[-1], all_distractors, list(reversed(true_ids_path)), get_random_entities(list(reversed(true_ids_path)), graph_algos.graph), graph_algos.graph)
        relevant_context_list = form_context_list(true_ids_path, graph_text_edge, graph_algos.graph, entity_top_alias=id2name)
        
        relevant_options_context_list = []
        for ent, parent_ent in options:
            relevant_text = graph_text_edge[parent_ent][ent]
            # if type(relevant_text) == list:
            #     relevant_text = ' '.join(relevant_text)
            relevant_options_context_list.append(relevant_text)
        
        all_context = get_all_context(true_ids_path + [ent for ent, _ in options], graph_text_sentencized)
        context_list = create_context_list(all_context, relevant_context_list, relevant_options_context_list, tokenizer, max_length=max_context_length)
        if len(context_list) == 0:
            print(f"path ids true: {true_ids_path}, query: {query}, options: {options}")
            raise ValueError("Tokenizer exceeded")
        if distractor is not None:
            ids_path.append(distractor)
            distractor_text = graph_text_edge[node_distracted][distractor]
            context_list.append(distractor_text)
        if shuffle_context:
            random.shuffle(context_list)
        context = '\n'.join([' '.join(context_part) for context_part in context_list])
        if id2name[true_ids_path[0]].lower() != entity_alias.lower():
            context += f" {id2name[true_ids_path[0]]} is also known as {entity_alias}."
        rel_path = [graph_algos.get_relation_for_vertex(true_ids_path[i], true_ids_path[i+1]) for i in range(len(true_ids_path)-1)]
        rel_info = [id2name[rel] for rel in rel_path]
        rel_aliases_used = query.split('->')[1:-1]
        # print(f"rel_info: {rel_info}, rel_aliases_used: {rel_aliases_used}")
        assert len(rel_info) == len(rel_aliases_used)
        rel_context = ' '.join([f"{rel_aliases_used[i]} means the same as {rel_info[i]}" for i in range(len(rel_info))])
        context += f"\n{rel_context}"
        answer_options = [ent for ent, _ in options]
        random.shuffle(answer_options)
        return {'query':query, 'correct_answers':[id2name[correct_id] for correct_id in correct_ids], 'path_id':true_ids_path, 
                'path_en':path, 'context':context, 'correct_ids':correct_ids, 'distractor':distractor, 'answer_options': answer_options, 
                'correct_ans_num': answer_options.index(true_ids_path[-1])+1}
### End of file: utils.py ###

### Start of file: server.py ###
# server.py

import socket
import json
import time
import argparse
from utils import *

checker_llm = None
def get_args():
    parser = argparse.ArgumentParser('Run Subgraph experiments')
    parser.add_argument('--checker_llm_device', type=str, default='cuda:1')

    return parser.parse_args()
def process_request(data):
    # Simulate processing a request by modifying the input data
    question = data['question']
    correct_answer = data['correct_answer']
    model_answer = data['model_answer']
    result = check_answer(question, correct_answer, model_answer)
    data['result'] = result
    return data

def check_answer(question, correct_answer, model_answer):
    global checker_llm
    return checker_llm.raw_checker(question=question, correct_ans=correct_answer, model_ans=model_answer)[0]

def main():
    host = 'localhost'
    port = 12345

    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
        s.bind((host, port))
        s.listen()

        print("Server is listening on", port)

        while True:
            conn, addr = s.accept()
            with conn:
                print('Connected by', addr)
                data = conn.recv(1024)
                if not data:
                    break
                # Process the request
                print('blocking...')
                try:
                    request_data = json.loads(data.decode('utf-8'))
                    response_data = process_request(request_data)
                    # Send response back to client
                    print('Sending response')
                    print(response_data)
                    conn.sendall(json.dumps(response_data).encode('utf-8'))
                except Exception as e:
                    print(e)
                    response_data = {}
                    response_data['result'] = 0
                    conn.sendall(json.dumps(response_data).encode('utf-8'))

if __name__ == '__main__':
    args = get_args()
    checker_llm = MistralChecker('mistralai/Mistral-7B-Instruct-v0.2', args.checker_llm_device)
    main()

### End of file: server.py ###

### Start of file: gpt_experiment.py ###
import numpy as np
import argparse
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
import torch
import argparse
import copy
from experiment_utils import *
import google.generativeai as genai
import time
from openai import OpenAI


BATCH_NUM = 1
qa_model = None
GPU_MAP = {0: "6GiB", 1: "8GiB", 2: "8GiB", 3: "8GiB", "cpu":"120GiB"}
INPUT_DEVICE = 'cuda:1'

def get_args():
    parser = get_base_args()
    parser.add_argument('--qa_llm', type=str, default='gpt-4o-mini')
    parser.add_argument('--quant_type', type=str, default=None, choices=[None])  # Explicitly set choices
    parser.set_defaults(num_queries=250) # override if needed
    return parser.parse_args()

def load_model(model_name="gpt-4o", only_tokenizer=False, gpu_map={0: "26GiB", 1: "0GiB", 2: "0GiB", 3: "0GiB", "cpu":"120GiB"}):
    tokenizer = None
    if not only_tokenizer:
        model = OpenAI()
        return tokenizer, model
    else:
        return tokenizer, None

def query_model(prompts, model, tokenizer, do_sample=True, top_k=10, 
                num_return_sequences=1, max_length=240, temperature=1.0, INPUT_DEVICE='cuda:0'):
    # preprocess prompts:
    responses = []
    for prompt in prompts:
        try:
            completion = model.chat.completions.create(
                model="gpt-4o-mini",
                messages=[
                    {"role": "system", "content": "You are a helpful assistant."},
                    {"role": "user", "content": f"{prompt}"},
                ],
                temperature=temperature,
                max_tokens=max_length,
                )
            responses.append(completion.choices[0].message.content)
        except Exception as e:
            #safety error
            print(e)
            response = "I'm sorry, I can't generate a response to that prompt."
            responses.append(response)
    time.sleep(0.01)
    return responses

def main():
    args = get_args()
    all_times, num_certificates_generated = run_experiment(args, load_model=load_model, query_model_func=query_model, 
                                                           GPU_MAP=GPU_MAP, BATCH_NUM=BATCH_NUM, INPUT_DEVICE=INPUT_DEVICE, model_context_length=7200)
    print(f'Completed {num_certificates_generated} certificates')
    print(f'Average time = {np.mean(all_times)}')
if __name__ == '__main__':
    main()
### End of file: gpt_experiment.py ###

### Start of file: phi_experiment.py ###
import numpy as np
import argparse
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
import torch
import argparse
import copy
from experiment_utils import *
import gc
import time

BATCH_NUM = 1
qa_model = None
GPU_MAP = {0: "40GiB", 1: "0GiB", 2: "0GiB", 3: "0GiB", "cpu":"120GiB"}
INPUT_DEVICE = 'cuda:0'
MODEL_CONTEXT_LEN = 120000
NUM_GEN = 0

def get_args():
    parser = get_base_args()
    parser.add_argument('--qa_llm', type=str, default='microsoft/Phi-3-mini-128k-instruct')
    parser.add_argument('--quant_type', type=str, default=None, choices=['8_bit', '4_bit'])  # Explicitly set choices
    parser.set_defaults(num_queries=250) # override if needed
    return parser.parse_args()

def load_model(model_name="microsoft/Phi-3-mini-128k-instruct", only_tokenizer=False, gpu_map={0: "26GiB", 1: "0GiB", 2: "0GiB", 3: "0GiB", "cpu":"120GiB"}, quant_type=None):
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    if not only_tokenizer:
        if quant_type is not None:
            if quant_type == '8_bit':
                print("loading 8 bit model")
                model = AutoModelForCausalLM.from_pretrained(model_name, device_map='auto', max_memory=gpu_map, torch_dtype=torch.float16, load_in_8bit=True, trust_remote_code=True, attn_implementation='flash_attention_2')
            elif quant_type == '4_bit':
                print("loading 4 bit model")
                model = AutoModelForCausalLM.from_pretrained(model_name, device_map='auto', max_memory=gpu_map, bnb_4bit_quant_type="nf4", load_in_4bit=True,  bnb_4bit_compute_dtype=torch.float16, trust_remote_code=True, attn_implementation='flash_attention_2')
        else:    
            model = AutoModelForCausalLM.from_pretrained(model_name, device_map='auto', max_memory=gpu_map, torch_dtype=torch.float16, trust_remote_code=True, attn_implementation='flash_attention_2')
        # model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16, trust_remote_code=True)
        for name, param in model.named_parameters():
            if 'cuda' not in str(param.device):
                print(f"param {name} not on cuda")
        # assert model.config.pad_token_id == tokenizer.pad_token_id, "The model's pad token ID does not match the tokenizer's pad token ID!"
        return tokenizer, model
    else:
        return tokenizer, None

def query_model(prompts, model, tokenizer, do_sample=True, top_k=10, 
                num_return_sequences=1, max_length=240, temperature=1.0, INPUT_DEVICE='cuda:0'):
    global NUM_GEN
    NUM_GEN += 1
    # preprocess prompts:
    start_time = time.time()
    PHI_SYS_PROMPT = "You are a helpful AI assistant. Answer questions as needed"
    chats = []
    if len(prompts) > 1:
        for prompt in prompts:
            message_template = [{"role": "system", "content": PHI_SYS_PROMPT}, {"role":"user", "content":f"{prompt}"}]
            chats.append([copy.deepcopy(message_template)])
    else:
        chats = [{"role": "system", "content": PHI_SYS_PROMPT}, {"role":"user", "content":f"{prompts[0]}"}]
        
    input_ids = tokenizer.apply_chat_template(chats, return_tensors="pt", add_generation_prompt=True, padding=True).to(INPUT_DEVICE)
    if input_ids.shape[-1] > 128000:
        print("Input too long, input too long, number of tokens: ", input_ids.shape)
        input_ids = input_ids[:, :128000]
    
    NUM_GEN += 1
    torch.cuda.empty_cache()
    if NUM_GEN % 50 == 0:
        gc.collect()

    start_t = time.time()
    generated_ids= model.generate(input_ids, max_new_tokens=max_length, do_sample=do_sample, temperature=temperature)
    print(f"Time taken for generating: {time.time() - start_t}")
    responses = tokenizer.batch_decode(generated_ids[:, input_ids.shape[-1]:].detach().cpu(), skip_special_tokens=True, clean_up_tokenization_spaces=True)
    
    del input_ids, generated_ids
    # print(responses)
    # torch.cuda.empty_cache()
    print(f"Time taken for generating: {time.time() - start_time}")
    return responses

def main():
    args = get_args()
    all_times, num_certificates_generated = run_experiment(args, load_model=load_model, query_model_func=query_model, 
                                                           GPU_MAP=GPU_MAP, BATCH_NUM=BATCH_NUM, INPUT_DEVICE=INPUT_DEVICE, model_context_length=MODEL_CONTEXT_LEN)
    print(f'Completed {num_certificates_generated} certificates')
    print(f'Average time = {np.mean(all_times)}')
if __name__ == '__main__':
    main()
### End of file: phi_experiment.py ###

### Start of file: gemini_experiment.py ###
import numpy as np
import argparse
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
import torch
import argparse
import copy
from experiment_utils import *
import google.generativeai as genai
import time

BATCH_NUM = 1
qa_model = None
GPU_MAP = {0: "6GiB", 1: "8GiB", 2: "8GiB", 3: "8GiB", "cpu":"120GiB"}
INPUT_DEVICE = 'cuda:1'
CONTINUOUS_SAFE = 0

def get_args():
    parser = get_base_args()
    parser.add_argument('--qa_llm', type=str, default='gemini-1.5-flash')
    parser.add_argument('--quant_type', type=str, default=None, choices=[None])  # Explicitly set choices
    parser.set_defaults(num_queries=250) # override if needed
    return parser.parse_args()

def load_model(model_name="gemini-1.5-flash", only_tokenizer=False, gpu_map={0: "26GiB", 1: "0GiB", 2: "0GiB", 3: "0GiB", "cpu":"120GiB"}, quant_type=None):
    global CONTINUOUS_SAFE
    tokenizer = None
    if not only_tokenizer:
        genai.configure(api_key=os.environ["API_KEY"])
        model = genai.GenerativeModel(model_name)
        return tokenizer, model
    else:
        return tokenizer, None

def query_model(prompts, model, tokenizer, do_sample=True, top_k=10, 
                num_return_sequences=1, max_length=120, temperature=1.0, INPUT_DEVICE='cuda:0'):
    # preprocess prompts:
    global CONTINUOUS_SAFE
    safe = [
        {
            "category": "HARM_CATEGORY_HARASSMENT",
            "threshold": "BLOCK_NONE",
        },
        {
            "category": "HARM_CATEGORY_HATE_SPEECH",
            "threshold": "BLOCK_NONE",
        },
        {
            "category": "HARM_CATEGORY_SEXUALLY_EXPLICIT",
            "threshold": "BLOCK_NONE",
        },
        {
            "category": "HARM_CATEGORY_DANGEROUS_CONTENT",
            "threshold": "BLOCK_NONE",
        },
    ]
    responses = []
    generation_config = genai.types.GenerationConfig(temperature=temperature, max_output_tokens=max_length)
    for prompt in prompts:
        try:
            response = model.generate_content(prompt, safety_settings=safe, generation_config=generation_config)
            responses.append(response.text)
            CONTINUOUS_SAFE = 0
        except Exception as e:
            #safety error
            print(e)
            response = "I'm sorry, I can't generate a response to that prompt."
            responses.append(response)
            time.sleep(1) #maybe rate limit error
            CONTINUOUS_SAFE += 1
            if CONTINUOUS_SAFE >= 4:
                print("Continuous safety error, too many", CONTINUOUS_SAFE, "possible rate limit error")
                exit(1)
    time.sleep(0.15)
    return responses

def main():
    args = get_args()
    all_times, num_certificates_generated = run_experiment(args, load_model=load_model, query_model_func=query_model, 
                                                           GPU_MAP=GPU_MAP, BATCH_NUM=BATCH_NUM, INPUT_DEVICE=INPUT_DEVICE, model_context_length=12800)
    print(f'Completed {num_certificates_generated} certificates')
    print(f'Average time = {np.mean(all_times)}')
if __name__ == '__main__':
    main()
### End of file: gemini_experiment.py ###

### Start of file: codebase2file.py ###
import os

def combine_python_files(source_dir, destination_file):
    """
    Combines the contents of all Python files (.py extension) in a source directory
    into a single text file.

    Args:
        source_dir: The path to the directory containing the Python files.
        destination_file: The path to the text file where the combined content will be written.
    """

    try:
        with open(destination_file, 'w', encoding='utf-8') as outfile:  # Open in write mode, handle encoding
            for filename in os.listdir(source_dir):
                source_path = os.path.join(source_dir, filename)

                if os.path.isfile(source_path) and filename.endswith(".py"):
                    try:
                        with open(source_path, 'r', encoding='utf-8') as infile:  # Open in read mode, handle encoding
                            outfile.write(f"### Start of file: {filename} ###\n") # File separator
                            outfile.write(infile.read())
                            outfile.write(f"\n### End of file: {filename} ###\n\n") # File separator
                        print(f"Added '{filename}' to '{destination_file}'")
                    except UnicodeDecodeError as e: # Handle UnicodeDecodeError if files have different encodings
                        print(f"Error reading '{filename}': {e}. Skipping this file.")

                    except IOError as e:
                        print(f"Unable to read '{filename}': {e}")

    except IOError as e:
        print(f"Unable to write to '{destination_file}': {e}")



if __name__ == "__main__":
    source_directory = input("Enter the source directory: ")
    destination_file_path = input("Enter the destination file path: ")

    combine_python_files(source_directory, destination_file_path)
    print("File combination complete.")
### End of file: codebase2file.py ###

### Start of file: experiment_utils.py ###
import os
import json
import numpy as np
from utils import *
import argparse
from statsmodels.stats.proportion import proportion_confint
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
import gc
from unidecode import unidecode
import pickle
import time
import torch
import gc
import socket
import gc
import argparse
import pickle
import copy

import argparse

def get_base_args():
    parser = argparse.ArgumentParser('Run Global experiments')
    parser.add_argument('--qa_graph_path', type=str, default='wikidata_graphs/wikidata_util.json')
    parser.add_argument('--context_graph_edge_path', type=str, default='wikidata_graphs/wikidata_text_edge.json')
    parser.add_argument('--results_dir', type=str, default='final_results/mistral_distractorfull/')
    parser.add_argument('--entity_aliases_path', type=str, default='wikidata5m_entity.txt')
    parser.add_argument('--id2name_path', type=str, default='wikidata_graphs/wikidata_name_id.json')
    parser.add_argument('--sentencized_path', type=str, default='wikidata_graphs/wikidata_sentencized.json')
    parser.add_argument('--relation_aliases_path', type=str, default='wikidata5m_relation.txt')
    parser.add_argument('--distractor_query', action='store_true', default=False, help=' best distractor based query?')
    parser.add_argument('--shuffle_context', action='store_true', default=False, help='Shuffle context in the context of query?')
    parser.add_argument('--k', type=int, default=4)
    parser.add_argument('--num_queries', type=int, default=250)
    parser.add_argument('--num_certificates', type=int, default=50)
    return parser

def experiment_pipeline(graph_algos, graph_text_edge, graph_text_sentencized, entity_aliases, source, relation_aliases, id2name, query_model, qa_model, tokenizer, model_context_length, k=5, distractor_query=False, num_queries=5, shuffle_context=True, BATCH_NUM=1, INPUT_DEVICE='cuda:0'):
    results = []
    correct = 0
    total = 0
    with torch.no_grad():
        for num_iter in range(num_queries//BATCH_NUM):
            prompts = []
            queries_data = []
            for j in range(BATCH_NUM):
                query_data = None
                while query_data is None:
                    query_data = get_query_data(graph_algos, source, id2name, graph_text_edge, graph_text_sentencized, tokenizer, 
                                                distractor_query=distractor_query, k=k, shuffle_context=shuffle_context, max_context_length=model_context_length)
                options_str = '\n'.join([f'{i+1}. {id2name[option]}' for i, option in enumerate(query_data['answer_options'])])
                prompt = LLM_PROMPT_TEMPLATE.format(context=query_data['context'], query=query_data['query'], options=options_str, few_shot_examples=FEW_SHOT_EXAMPLES)
                prompts.append(prompt)
                queries_data.append(query_data)
                if distractor_query:
                    assert len(query_data['path_id']) >= 2
            model_answers= query_model(prompts, qa_model, tokenizer, temperature=0.000001, INPUT_DEVICE=INPUT_DEVICE, do_sample=False) #do_sample=False, so no sampling or temp=0
            for i, model_ans in enumerate(model_answers):
                model_ans = model_ans.strip()
                model_answers[i] = model_ans
            assert len(queries_data) == len(model_answers)
            for i in range(len(queries_data)):
                query = queries_data[i]['query']
                correct_answers = queries_data[i]['correct_answers']
                path = queries_data[i]['path_en']
                path_id = queries_data[i]['path_id']
                context = queries_data[i]['context']
                correct_ids = queries_data[i]['correct_ids']
                distractor = queries_data[i]['distractor']
                model_ans = model_answers[i]
                if len(model_ans) == 0:
                    continue
                eval_ans = 0
                for num_correct, correct_answer in enumerate(correct_answers):
                    eval_ans = dumb_checker(model_ans, queries_data[i]['correct_ans_num'])
                    assert id2name[correct_ids[num_correct]] == correct_answer
                    if eval_ans == 1:
                        break
                results.append({ 'question':query, 'correct_answers':correct_answers, 'model_answer':model_ans, 
                                'path_en':path, 'path_id':path_id, 'context':context, 'result':(eval_ans, None), 
                                'distractor':distractor, 'correct_ids':correct_ids, 'options':queries_data[i]['answer_options'], 'correct_ans_num':queries_data[i]['correct_ans_num']})
                correct += results[-1]['result'][0]
                total += 1
            print(f'Completed {num_iter+1} queries, {correct} correct out of {total} total')
            del model_answers
        print(f'Completed {num_queries} queries, {correct} correct out of {total} total')
    return results, correct, total

def load_experiment_setup(args, load_model, GPU_MAP):
    if not os.path.exists(args.results_dir):
        os.makedirs(args.results_dir)
    tokenizer, qa_model = load_model(args.qa_llm, only_tokenizer=False, gpu_map=GPU_MAP, quant_type=args.quant_type)
    qa_graph = json.load(open(args.qa_graph_path))
    context_graph_edge = json.load(open(args.context_graph_edge_path))
    graph_text_sentencized = json.load(open(args.sentencized_path))
    id2name = json.load(open(args.id2name_path))
    entity_aliases = load_aliases(args.entity_aliases_path)
    relation_aliases = load_aliases(args.relation_aliases_path)
    print(f"Best Distractor Task: {args.distractor_query}")
    qa_graph_algos = GraphAlgos(qa_graph, entity_aliases, relation_aliases)
    # best_vertices = qa_graph_algos.get_best_vertices(num=1000)
    best_vertices =  ['Q38', 'Q1055', 'Q838292', 'Q34433', 'Q254', 'Q31', 'Q270', 'Q200482', 'Q36740', 'Q1911276', 'Q3740786', 'Q1124384', 'Q931739', 'Q2090699', 'Q505788', 'Q1217787', 'Q115448', 'Q2502106', 'Q1793865', 'Q229808', 'Q974437', 'Q219776', 'Q271830', 'Q279164', 'Q76508', 'Q245392', 'Q2546120', 'Q312408', 'Q6110803', 'Q211196', 'Q18407657', 'Q18602670', 'Q21979809', 'Q23010088', 'Q1338555', 'Q5516100', 'Q1765358', 'Q105624', 'Q166262', 'Q33', 'Q36', 'Q16', 'Q96', 'Q36687', 'Q282995', 'Q858401', 'Q850087', 'Q864534', 'Q291244', 'Q159', 'Q668', 'Q211', 'Q183', 'Q1603', 'Q408', 'Q218'][:50]
    # random.shuffle(best_vertices)
    if not os.path.exists(args.results_dir):
        os.makedirs(args.results_dir, exist_ok=True)
    
    return qa_graph_algos, context_graph_edge, graph_text_sentencized, entity_aliases, relation_aliases, id2name, qa_model, tokenizer, best_vertices

def run_experiment(args, load_model, query_model_func, GPU_MAP, model_context_length, BATCH_NUM=1, INPUT_DEVICE='cuda:0'):
    qa_graph_algos, context_graph_edge, graph_text_sentencized, entity_aliases, relation_aliases, id2name, qa_model, tokenizer, best_vertices = load_experiment_setup(args, load_model, GPU_MAP)
    already_done = []
    if not os.path.exists(args.results_dir):
        os.makedirs(args.results_dir, exist_ok=True)
    
    for file in os.listdir(args.results_dir):
        idx = file.index('Q')
        vertex_id = file[idx:-4]
        already_done.append(vertex_id)
    all_times = []
    num_certificates_generated = 0
    for i, vertex_id in enumerate(best_vertices):
        if num_certificates_generated >= args.num_certificates:
            break
        vertex = id2name[vertex_id]
        if vertex_id in already_done:
            print("Already done", vertex_id, vertex)
            continue
        start_time = time.time()
        subgraph = qa_graph_algos.create_subgraph_within_radius(vertex_id, 4)
        subgraph_algos = GraphAlgos(subgraph, entity_aliases, relation_aliases)
        if len(subgraph) < 900:
            print(len(subgraph), "Skipping", vertex_id, vertex) # Skip small subgraphs
            continue
        print(vertex, vertex_id, len(subgraph))
        num_certificates_generated += 1
        
        experiment_results = experiment_pipeline(graph_algos=subgraph_algos, graph_text_edge=context_graph_edge, graph_text_sentencized=graph_text_sentencized, 
                                                 entity_aliases=entity_aliases, source=vertex_id, relation_aliases=relation_aliases, id2name=id2name, 
                                                 query_model=query_model_func, qa_model=qa_model, tokenizer=tokenizer, k=args.k, 
                                                 distractor_query=args.distractor_query, num_queries=args.num_queries, 
                                                 shuffle_context=args.shuffle_context, BATCH_NUM=BATCH_NUM, INPUT_DEVICE=INPUT_DEVICE, model_context_length=model_context_length)
        end_time = time.time()
        print(f'Time taken for {vertex} = {end_time - start_time}')
        all_times.append(end_time - start_time)
        with open(os.path.join(args.results_dir, str(vertex_id)+'.pkl'), 'wb') as f:
            pickle.dump(experiment_results, f)
    return all_times, num_certificates_generated
### End of file: experiment_utils.py ###

### Start of file: mistral_experiment.py ###
import numpy as np
import argparse
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
import torch
import argparse
import copy
from experiment_utils import *

BATCH_NUM = 1
qa_model = None
GPU_MAP = {0: "0GiB", 1: "20GiB", 2: "20GiB", 3: "0GiB", "cpu":"0GiB"}
INPUT_DEVICE = 'cuda:1'
MAX_CONTEXT_LEN = 28000
NUM_GEN = 0

def get_args():
    parser = get_base_args()
    parser.add_argument('--qa_llm', type=str, default='mistralai/Mistral-7B-Instruct-v0.2')
    parser.add_argument('--quant_type', type=str, default=None, choices=['8_bit', '4_bit'])  # Explicitly set choices
    parser.set_defaults(num_queries=250) # override if needed
    return parser.parse_args()

def load_model(model_name="mistralai/Mistral-7B-Instruct-v0.2", only_tokenizer=False, gpu_map={0: "26GiB", 1: "0GiB", 2: "0GiB", 3: "0GiB", "cpu":"120GiB"}, quant_type=None):
    tokenizer = AutoTokenizer.from_pretrained(model_name, padding_side='left')
    if not only_tokenizer:
        if quant_type is not None:
            if quant_type == '8_bit':
                print("loading 8 bit model")
                model = AutoModelForCausalLM.from_pretrained(model_name, device_map='auto', max_memory=gpu_map, torch_dtype=torch.float16, load_in_8bit=True)
            elif quant_type == '4_bit':
                print("loading 4 bit model")
                model = AutoModelForCausalLM.from_pretrained(model_name, device_map='auto', max_memory=gpu_map, bnb_4bit_quant_type="nf4", load_in_4bit=True,  bnb_4bit_compute_dtype=torch.float16)
        else:
            print('no quantization, loading in fp16')
            model = AutoModelForCausalLM.from_pretrained(model_name, device_map='auto', max_memory=gpu_map, torch_dtype=torch.float16)
        #check device of all model tensors
        for name, param in model.named_parameters():
            if 'cuda' not in str(param.device):
                print(f"param {name} not on cuda")
        return tokenizer, model
    else:
        return tokenizer, None

def query_model(prompts, qa_model, tokenizer, do_sample=True, top_k=10, 
                num_return_sequences=1, max_length=240, temperature=1.0, INPUT_DEVICE='cuda:0'):
    global NUM_GEN
    # preprocess prompts:
    import time
    assert len(prompts) == 1
    
    messages = [{"role": "user", "content": f"{prompts[0]}"},]

    encodeds = tokenizer.apply_chat_template(messages, return_tensors="pt")
    model_inputs = encodeds.to(INPUT_DEVICE)
    start_time = time.time()
    generated_ids = qa_model.generate(model_inputs, max_new_tokens=max_length, do_sample=do_sample, temperature=temperature)
    print(f"Time taken for model: {time.time() - start_time}")
    generated_ids = generated_ids[:, model_inputs.shape[-1]:]
    decoded = tokenizer.batch_decode(generated_ids.detach().cpu())

    NUM_GEN += 1
    if NUM_GEN % 50 == 0:
        torch.cuda.empty_cache()
        gc.collect()
    model_ans = decoded[0].strip()
    del model_inputs, generated_ids
    return [model_ans]

def main():
    args = get_args()
    all_times, num_certificates_generated = run_experiment(args, load_model=load_model, query_model_func=query_model, 
                                                           GPU_MAP=GPU_MAP, BATCH_NUM=BATCH_NUM, INPUT_DEVICE=INPUT_DEVICE, model_context_length=MAX_CONTEXT_LEN)
    print(f'Completed {num_certificates_generated} certificates')
    print(f'Average time = {np.mean(all_times)}')
if __name__ == '__main__':
    main()
### End of file: mistral_experiment.py ###

### Start of file: llama_experiment.py ###
import numpy as np
import argparse
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
import torch
import argparse
import copy
from experiment_utils import *
from transformers import BitsAndBytesConfig

BATCH_NUM = 1
qa_model = None
GPU_MAP = {0: "0GiB", 1: "0GiB", 2: "40GiB", 3: "40GiB", "cpu":"50GiB"}
INPUT_DEVICE = 'cuda:2'

def get_args():
    parser = get_base_args()
    parser.add_argument('--qa_llm', type=str, default='meta-llama/Meta-Llama-3-8B-Instruct')
    parser.add_argument('--quant_type', type=str, default=None, choices=['8_bit', '4_bit'])  # Explicitly set choices
    parser.set_defaults(num_queries=250) # override if needed
    return parser.parse_args()

def load_model(model_name="lmsys/vicuna-13b-v1.5", only_tokenizer=False, gpu_map={0: "26GiB", 1: "0GiB", 2: "0GiB", 3: "0GiB", "cpu":"120GiB"}, quant_type=None):
    tokenizer = AutoTokenizer.from_pretrained(model_name, padding_side='left')
    if not only_tokenizer:
        if quant_type is not None:
            if quant_type == '8_bit':
                print("loading 8 bit model")
                model = AutoModelForCausalLM.from_pretrained(model_name, device_map='auto', max_memory=gpu_map, torch_dtype=torch.float16, load_in_8bit=True)
            elif quant_type == '4_bit':
                print("loading 4 bit model")
                model = AutoModelForCausalLM.from_pretrained(model_name, device_map='auto', max_memory=gpu_map, bnb_4bit_quant_type="nf4", load_in_4bit=True,  bnb_4bit_compute_dtype=torch.float16)
        else:
            print('no quantization, loading in fp16')
            model = AutoModelForCausalLM.from_pretrained(model_name, device_map='auto', max_memory=gpu_map, torch_dtype=torch.float16)
        tokenizer.pad_token = tokenizer.eos_token
        model.config.pad_token_id = tokenizer.pad_token_id
        assert model.config.pad_token_id == tokenizer.pad_token_id, "The model's pad token ID does not match the tokenizer's pad token ID!"
        tokenizer.padding_side = 'right'
        return tokenizer, model
    else:
        return tokenizer, None

def query_model(prompts, model, tokenizer, do_sample=True, top_k=10, 
                num_return_sequences=1, max_length=240, temperature=1.0, INPUT_DEVICE='cuda:0'):
    # preprocess prompts:
    LLAMA3_SYS_PROMPT = "You are a helpful chatbot who answers multiple choice reasoning questions in a specified format choosing from only the options available"
    chats = []
    if len(prompts) > 1:
        for prompt in prompts:
            message_template = [{"role": "system", "content": LLAMA3_SYS_PROMPT}, {"role":"user", "content":f"{prompt}"}]
            chats.append([copy.deepcopy(message_template)])
    else:
        chats = [{"role": "system", "content": LLAMA3_SYS_PROMPT}, {"role":"user", "content":f"{prompts[0]}"}]
        
    input_ids = tokenizer.apply_chat_template(chats, return_tensors="pt", add_generation_prompt=True, padding=True).to(INPUT_DEVICE)
    if input_ids.shape[-1] > 8000:
        print("Input too long, input too long, number of tokens: ", input_ids.shape)
        input_ids = input_ids[:, :8000]
    terminators = [
        tokenizer.eos_token_id,
        tokenizer.convert_tokens_to_ids("<|eot_id|>")
    ]
    generated_ids= model.generate(input_ids, max_new_tokens=max_length, do_sample=do_sample, eos_token_id=terminators, temperature=temperature)
    responses = tokenizer.batch_decode(generated_ids[:, input_ids.shape[-1]:].detach().cpu(), skip_special_tokens=True, clean_up_tokenization_spaces=True)
    
    del input_ids, generated_ids

    return responses

def main():
    args = get_args()
    all_times, num_certificates_generated = run_experiment(args, load_model=load_model, query_model_func=query_model, 
                                                           GPU_MAP=GPU_MAP, BATCH_NUM=BATCH_NUM, INPUT_DEVICE=INPUT_DEVICE, model_context_length=7200)
    print(f'Completed {num_certificates_generated} certificates')
    print(f'Average time = {np.mean(all_times)}')
if __name__ == '__main__':
    main()
### End of file: llama_experiment.py ###

### Start of file: wikidata_make.py ###
import os
import json
import gc
from unidecode import unidecode
import nltk
nltk.download('punkt')
from nltk.tokenize import sent_tokenize
import re
import time
import copy
from utils import load_aliases

wikidata_graph = {}
wikidata_name_id = {}
wikidata_en_graph = {}
wikidata_en_graph = {}
wikidata_text = {}
wikidata_en_text = {}
codex_graph = {}

MIN_CONTEXT_LEN = 4
MAX_CONTEXT_LEN = 10

find = ['P31', 'P279', 'P361', 'P2184', 'P2633'] # instance of, subclass of, part of, these relations are very vague so we discard them
with open('wikidata5m_all_triplet.txt', 'r') as file: #TODO: fix hard-coded paths; where's the file in the repo?
    for line in file:
        line = line.strip()
        line = line.split('\t')
        line = [x.strip() for x in line]
        if line[1] in find:
            continue
        if line[0] not in wikidata_graph:
            wikidata_graph[line[0]] = {}
        if line[1] not in wikidata_graph[line[0]]:
            wikidata_graph[line[0]][line[1]] = []
        wikidata_graph[line[0]][line[1]].append(line[2])
    
with open('wikidata5m_text.txt', 'r', encoding="utf-8") as f: # where is this file?
    for line in f:
        line = line.strip()
        line = line.split('\t')
        line = [x.strip() for x in line]
        text = line[1].strip()
        if len(line) > 2:
            for i in range(2, len(line)):
                text += ' ' + line[i].strip()
        if line[0] in wikidata_graph:
            wikidata_text[line[0]] = text


possible_entity_names = load_aliases('wikidata5m_entity.txt')

#TODO: discuss the most common alias, do we continue filtering out nodes where abstract do not mention nodes
# We primarily use this for two things: ensure some alias occurs in node's text, also the sentence in query of the form : A is also as A_alias uses info from this for naming A
print("Starting Names")
wikidata_name_id = {}
with open('wikidata5m_entity.txt', 'r') as f: # where is this file?
    for line in f:
        line = line.strip()
        line = line.split('\t')
        line = [x.strip() for x in line]
        if line[0] not in wikidata_text:
            wikidata_name_id[line[0]] = unidecode(possible_entity_names[line[0]][0]).lower()
            continue
        possible_entity_names[line[0]] = line[1:]
        possible = line[1:] #select most common label from the aliases
        common = ''
        all_text = wikidata_text[line[0]]
        all_text = unidecode(all_text)
        all_text = all_text.lower()
        all_text = all_text.replace('–', '-') #replace en dash with hyphen
        min_ind = 1000000
        #choose common according to the order of the aliases appearing in the text
        for name in possible:
            name = unidecode(name)
            name = name.lower()
            pattern = r'\b' + re.escape(name) + r'\b'
            matches = re.search(pattern, all_text)
            if matches:
                ind = all_text.index(name)
                if ind < min_ind:
                    min_ind = ind
                    common = name
                if min_ind <= 15:
                    #index of alias is early enough in text to break and not search for an alias occuring earlier
                    break
        if common == '':
            common = unidecode(possible_entity_names[line[0]][0]).lower()
        wikidata_name_id[line[0]] = common.strip()
possible_relation_names = {}
with open('wikidata5m_relation.txt', 'r') as f: # where is this file?
    for line in f:
        line = line.strip()
        line = line.split('\t')
        line = [x.strip() for x in line]
        wikidata_name_id[line[0]] = unidecode(line[1]).lower()
        possible_relation_names[line[0]] = line[1:]
gc.collect() # is this needed? 

print("len wikidata name id:", len(wikidata_name_id))

#add death place birth place sentences
relid2sent = {'P20':'died in', 'P19':'was born in'}
for key in wikidata_graph:
    if key not in wikidata_text:
        continue
    key_text = ''
    if key not in wikidata_name_id:
        if key not in possible_entity_names:
            continue
        key_text = possible_entity_names[key][0]
    else:
        key_text = wikidata_name_id[key]
    for relid in relid2sent:
        if relid in wikidata_graph[key]:
            if wikidata_graph[key][relid][0] in wikidata_name_id:
                wikidata_text[key] += f'{key_text} {relid2sent[relid]} {wikidata_name_id[wikidata_graph[key][relid][0]]}.'
            elif wikidata_graph[key][relid][0] not in possible_entity_names:
                continue
            else:
                wikidata_text[key] += f'{key_text} {relid2sent[relid]} {possible_entity_names[wikidata_graph[key][relid][0]][0]}.'
                
wikidata_en_graph = {}
wikidata_id_name = {}
for key, value in wikidata_name_id.items():
    wikidata_id_name[value] = key
    
wikidata_util = {}
for key, value in wikidata_graph.items():
    wikidata_util[key] = {}
    for rel, objs  in value.items():
        for obj in objs:
            wikidata_util[key][obj] = rel

print(len(wikidata_util), len(wikidata_graph))

wikidata_text_edge = {}


wikidata_text_sentencized = {}
for key, value in wikidata_text.items():
    sentencized = sent_tokenize(value)
    if len(sentencized) == 0:
        print(key, value)
        raise ValueError('No sentence found')
    sentencized = [x.strip() for x in sentencized]
    wikidata_text_sentencized[key] = sentencized

wikidata_text_sentencized_format = {}
for key, value in wikidata_text_sentencized.items():
    wikidata_text_sentencized_format[key] = [unidecode(sent).lower().strip() for sent in value]
len(wikidata_text_sentencized_format), len(wikidata_text_sentencized)

print("Starting Context Trimming")
new_wikidata = {}
wikidata_text_edge = {}
count = 0
total = 0
keys_done_num = 0
keys_considered = 0
start_time = time.time()
for key, value in wikidata_util.items():
    keys_considered += 1
    if key not in wikidata_text or key not in wikidata_name_id:
        continue
    keys_done_num += 1
    wikidata_text_edge[key] = {}
    possible_key_names = possible_entity_names[key]
    possible_key_names = [unidecode(x).lower().strip() for x in possible_key_names]
    new_value = {}
    for obj, rel in value.items():
        total += 1
        if obj not in wikidata_name_id or rel not in wikidata_name_id:
            continue
        possible_names = possible_entity_names[obj]
        possible_names = [unidecode(x).lower().strip() for x in possible_names]
        relation_names = possible_relation_names[rel]
        relation_names = [unidecode(x).lower().strip() for x in relation_names]
        key_relevant_ids = []
        obj_relevant_ids = []
        key_relevant_ids.append(0)
        # relevant_sentences.append(wikidata_text_sentencized[key][0])
        found = False
                
        for i, sent in enumerate(wikidata_text_sentencized_format[key]):
            for name in possible_names:
                if len(name) <= 0:
                    continue
                pattern = r'\b' + re.escape(name) + r'\b'
                matches = re.search(pattern, sent)
                if matches:
                    if i > 0:
                        key_relevant_ids.append(i)
                    new_value[obj] = rel
                    found = True
                    break
        if not found:
            if obj not in wikidata_text:
                continue
            obj_relevant_ids.append(0)
            # relevant_sentences.append(wikidata_text_sentencized[obj][0])
            for i, sent in enumerate(wikidata_text_sentencized_format[obj]):
                for name in possible_key_names:
                    if len(name) <= 0:
                        continue
                    pattern = r'\b' + re.escape(name) + r'\b'
                    matches = re.search(pattern, sent)
                    if matches:
                        # print(i, len(wikidata_text_sentencized_format[obj]))
                        if i > 0:
                            obj_relevant_ids.append(i)
                        new_value[obj] = rel
                        found = True
                        break

        if found:
            for i, sent in enumerate(wikidata_text_sentencized_format[key]):
                if i not in key_relevant_ids:
                    for name in relation_names:
                        if len(name) <= 0:
                            continue
                        pattern = r'\b' + re.escape(name) + r'\b'
                        matches = re.search(pattern, sent)
                        if matches:
                            if i > 0:
                                key_relevant_ids.append(i)
                            new_value[obj] = rel
                            found = True
                            break
            key_relevant_ids = list(sorted(key_relevant_ids))
            relevant_sentences = []
            num_key_sents = len(wikidata_text_sentencized[key])
            k_done = 0 
            k_todo = MIN_CONTEXT_LEN - len(key_relevant_ids) - len(obj_relevant_ids) # additional sentences to be added as 5 is min in context
            
            #add all key sentences as well as filler sentences in between to get min length context
            for i in range(len(key_relevant_ids)):
                relevant_sentences.append(wikidata_text_sentencized[key][key_relevant_ids[i]])
                if k_done >= k_todo:
                    continue
                if i < len(key_relevant_ids) - 1:
                    if key_relevant_ids[i+1] - key_relevant_ids[i] > 1:
                        #add atmost 2 sentences between key sentences if needed
                        for j in range(max(key_relevant_ids[i]+1, key_relevant_ids[i+1]-2), key_relevant_ids[i+1]):
                            if k_done >= k_todo:
                                break
                            relevant_sentences.append(wikidata_text_sentencized[key][j])
                            k_done += 1
                else:
                    #add atmost 2 sentences after the last key sentence if needed
                    if key_relevant_ids[i] < num_key_sents - 1:
                        for j in range(max(key_relevant_ids[i]+1, num_key_sents-2), num_key_sents):
                            if k_done >= k_todo:
                                break
                            relevant_sentences.append(wikidata_text_sentencized[key][j])
                            k_done += 1
            
            #max context sentences number is 8, #TODO: we could change this for models with large context
            if len(relevant_sentences) > MAX_CONTEXT_LEN:
                relevant_sentences = relevant_sentences[:MAX_CONTEXT_LEN]
            #as we add and remove in order the intro sentence about the entity is always present
            assert wikidata_text_sentencized[key][0] in relevant_sentences 
            for idx in obj_relevant_ids:
                if len(relevant_sentences) >= MAX_CONTEXT_LEN:
                    break
                relevant_sentences.append(wikidata_text_sentencized[obj][idx])
            wikidata_text_edge[key][obj] = copy.deepcopy(relevant_sentences) #' '.join(relevant_sentences)
            new_value[obj] = rel
            if count % 1000000 == 0:
                print(key, obj, len(wikidata_text_edge[key][obj]), wikidata_text_edge[key][obj])
            count += 1
    if len(new_value) > 0:
        new_wikidata[key] = new_value
    if keys_done_num % 100000 == 0:
        end_time = time.time()
        print("Time taken for 100000 keys:", end_time - start_time)
        start_time = time.time()
        print(keys_done_num, keys_considered, count, total)

print("Total keys:", keys_done_num, keys_considered, count, total)

remove = set()
new_graph = {}
for key, value in new_wikidata.items():
    new_value = {}
    for k, v in value.items():
        if v not in wikidata_name_id:
            remove.add(v)
            continue
        new_value[k] = v
    if len(new_value) > 0:
        new_graph[key] = new_value

wikidata_graph_util = new_graph
os.makedirs('wikidata_graphs', exist_ok=True)
with open('wikidata_graphs/wikidata_name_id.json', 'w') as f:
    json.dump(wikidata_name_id, f)
with open('wikidata_graphs/wikidata_text.json', 'w') as f:
    json.dump(wikidata_text, f)
with open('wikidata_graphs/wikidata_util.json', 'w') as f:
    json.dump(wikidata_graph_util, f)
with open('wikidata_graphs/wikidata_sentencized.json', 'w') as f:
    json.dump(wikidata_text_sentencized, f)
with open('wikidata_graphs/wikidata_text_edge.json', 'w') as f:
    json.dump(wikidata_text_edge, f)

for key, value in wikidata_graph_util.items():
    for k, v in value.items():
        assert k in wikidata_text_edge[key]

os.makedirs('wikidata_graphs', exist_ok=True)
with open('wikidata_graphs/wikidata_name_id.json', 'w') as f:
    json.dump(wikidata_name_id, f)
with open('wikidata_graphs/wikidata_text.json', 'w') as f:
    json.dump(wikidata_text, f)
with open('wikidata_graphs/wikidata_util.json', 'w') as f:
    json.dump(wikidata_graph_util, f)
with open('wikidata_graphs/wikidata_sentencized.json', 'w') as f:
    json.dump(wikidata_text_sentencized, f)
with open('wikidata_graphs/wikidata_text_edge.json', 'w') as f:
    json.dump(wikidata_text_edge, f)
### End of file: wikidata_make.py ###