icl.py

import argparse
import csv
import copy
import random
import json
import os
import re
import numpy as np
import string
import pickle



_MLN = '''
The reasoning process of a Markov Logic Network (MLN) is:
1. Represent Knowledge: Use weighted first-order logic formulas, where higher weights indicate stronger rules.
2. Ground Formulas: For a given set of predicates, create specific instances (or worlds) to evaluate the truth of these formulas.
3. Calculate Probability: For each possible world (x), compute the score as follows:
   score(x) = \sum_i w_i \cdot n_i(x)
   where (w_i) are the weights and n_i(x) is the number of true groundings of formula (i) in world (x).
4. Normalize Scores: Calculate the probability of a world (x) using the formula:
   P(X = x) = \exp(score(x)) / Z
   where Z = \sum_{x'} \exp{score(x')} is the partition function that normalizes the probabilities across all possible worlds.
By following these steps, you will emulate the MLN's ability to reason about uncertain knowledge in structured environments.
'''

_TASK_DESCRIPTION_COT = '''
Read how markolv logic inference works:
{mln}
Imagine you are a Markov Logic Network engine. \
Your task is to perform probabilistic reasoning based on a combination of rules and known facts.
Provide the most probable answer for the question. Please ensure that you are based solely on the information provided in the prompt, rather than relying on your existing internal knowledge. Solve step by step to show how the answer is derived. Wrap the final answer in <answer> and </answer> tags
'''

_TASK_DESCRIPTION = '''
Read how markolv logic inference works:
{mln}
Imagine you are a Markov Logic Network engine. \
Your task is to perform probabilistic reasoning based on a combination of rules and known facts.
Provide the most probable answer for the question. Please ensure that you are based solely on the information provided in the prompt, rather than relying on your existing internal knowledge. 
Write the answer in the following format {{answer is: <ANSWER>}}
'''

class Baseline():

    def __init__(self, prompt_mode, tokenizer, cot_mode=None, has_system_prompt=True):
        self.prompt_mode = prompt_mode
        self.tokenizer = tokenizer
        self.task_descript = _TASK_DESCRIPTION.format(mln=_MLN)
        self.has_system_prompt = has_system_prompt
        if self.prompt_mode == 'icl':
            if not has_system_prompt:
                self.msges = [{"role": "user", "content": _TASK_DESCRIPTION.format(mln=_MLN)}]    
            else:
                self.msges = [{"role": "system", "content": _TASK_DESCRIPTION.format(mln=_MLN)}]
        elif self.prompt_mode == 'cot':
            if not has_system_prompt:
                self.msges = [{"role": "user", "content": _TASK_DESCRIPTION_COT.format(mln=_MLN)}]    
            else:
                self.msges = [{"role": "system", "content": _TASK_DESCRIPTION_COT.format(mln=_MLN)}]
        else:
            self.msges = []
        self.cot_mode = cot_mode

    def make_icl_demo(self, inp, first_demo=False):


        user = ""
        user += inp['queries'][0] + "\n"
        return [
                {"role": "user", "content": user.strip()},
                {"role": "assistant", "content": "{answer is: "+inp['queries'][1].strip()+"}"}
            ]

    def make_query_prompt_fix(self, inp, ret='str'):
        user = ""
        user += inp['queries'][0] + "\n"


        msges = self.msges + [{'role': "user", "content": user.strip()}]
        ret = self.tokenizer.apply_chat_template(msges, tokenize=False, add_generation_prompt=True)

        return ret
        

    def proc_demos(self, demos):
        for x in demos:
            if self.prompt_mode == "icl":
                self.msges += self.make_icl_demo(x)
            else:
                continue
                #self.msges += self.make_cot_demo(x)
    



class HFPrompterICL():

    def __init__(self, prompt_mode, tokenizer, cot_mode=None, has_system_prompt=True):
        self.prompt_mode = prompt_mode
        self.tokenizer = tokenizer
        self.task_descript = _TASK_DESCRIPTION.format(mln=_MLN)
        self.has_system_prompt = has_system_prompt
        if self.prompt_mode == 'icl':
            if not has_system_prompt:
                self.msges = [{"role": "user", "content": _TASK_DESCRIPTION.format(mln=_MLN)}]    
            else:
                self.msges = [{"role": "system", "content": _TASK_DESCRIPTION.format(mln=_MLN)}]
        elif self.prompt_mode == 'cot':
            if not has_system_prompt:
                self.msges = [{"role": "user", "content": _TASK_DESCRIPTION_COT.format(mln=_MLN)}]    
            else:
                self.msges = [{"role": "system", "content": _TASK_DESCRIPTION_COT.format(mln=_MLN)}]
        else:
            self.msges = []
        self.cot_mode = cot_mode

    def make_icl_demo(self, inp, first_demo=False):


        user = ""
        user += "Rules and their weights:\n"
        if not 'rweight' in inp.keys():
            for r in inp['rules']:
                user += str(1) + " " + r + '\n'
        else:
            for r, w in zip(inp['rules'], inp['rweight']):
                user += str(w) + " " + r + '\n'

        qa_pairs = '\n'.join([x[0] for x in inp['facts']])

        user += f"Known facts:\n {qa_pairs}"+"\n"
        
        user += inp['queries'][0] + "\n"

        return [
                {"role": "user", "content": user.strip()},
                {"role": "assistant", "content": "{answer is: "+inp['queries'][1].strip()+"}"}
            ]

    def make_query_prompt_fix(self, inp, ret='str'):
                  
        user = "Rules and their weights:\n"
        if not 'rweight' in inp.keys():
            for r in inp['rules']:
                user += str(1) + " " + r + '\n'
        else:
            for r, w in zip(inp['rules'], inp['rweight']):
                user += str(w) + " " + r + '\n'

        qa_pairs = '\n'.join([x[0] for x in inp['facts']])
        user += f"Known facts:\n {qa_pairs}"+"\n"
        user += inp['queries'][0] + "\n"
        if self.cot_mode == 'cot':
            user += "Let's think step by step.\n"

        msges = self.msges + [{'role': "user", "content": user.strip()}]
        
        ret = self.tokenizer.apply_chat_template(msges, tokenize=False, add_generation_prompt=True)
        print(ret)
        return ret
        

    def proc_demos(self, demos):
        for x in demos:
            if self.prompt_mode == "icl":
                self.msges += self.make_icl_demo(x)
            else:
                continue
                #self.msges += self.make_cot_demo(x)
    

class OpenaiICL():

    def __init__(self, prompt_mode, cot_mode=None, has_system_prompt=True):
        self.prompt_mode = prompt_mode
        self.task_descript = _TASK_DESCRIPTION.format(mln=_MLN)
        self.has_system_prompt = has_system_prompt
        if self.prompt_mode == 'icl':
            if not has_system_prompt:
                self.msges = [{"role": "user", "content": _TASK_DESCRIPTION.format(mln=_MLN)}]    
            else:
                self.msges = [{"role": "system", "content": _TASK_DESCRIPTION.format(mln=_MLN)}]
        elif self.prompt_mode == 'cot':
            if not has_system_prompt:
                self.msges = [{"role": "user", "content": _TASK_DESCRIPTION_COT.format(mln=_MLN)}]    
            else:
                self.msges = [{"role": "system", "content": _TASK_DESCRIPTION_COT.format(mln=_MLN)}]
        else:
            self.msges = []
        self.cot_mode = cot_mode

    def make_icl_demo(self, inp, first_demo=False):
        user = ""
        user += "Rules and their weights:\n"
        if not 'rweight' in inp.keys():
            for r in inp['rules']:
                user += str(1) + " " + r + '\n'
        else:
            for r, w in zip(inp['rules'], inp['rweight']):
                user += str(w) + " " + r + '\n'

        qa_pairs = '\n'.join([x[0] for x in inp['facts']])

        user += f"Known facts:\n {qa_pairs}"+"\n"
        
        user += inp['queries'][0] + "\n"

        return [
                {"role": "user", "content": user.strip()},
                {"role": "assistant", "content": "{answer is: "+inp['queries'][1].strip()+"}"}
            ]

    def make_query_prompt_fix(self, inp, ret='str'):
                  
        user = "Rules and their weights:\n"
        if not 'rweight' in inp.keys():
            for r in inp['rules']:
                user += str(1) + " " + r + '\n'
        else:
            for r, w in zip(inp['rules'], inp['rweight']):
                user += str(w) + " " + r + '\n'

        qa_pairs = '\n'.join([x[0] for x in inp['facts']])
        user += f"Known facts:\n {qa_pairs}"+"\n"
        user += inp['queries'][0] + "\n"
        if self.cot_mode == 'cot':
            user += "Let's think step by step.\n"

        msges = self.msges + [{'role': "user", "content": user.strip()}]

        return msges
        

    def proc_demos(self, demos):
        for x in demos:
            if self.prompt_mode == "icl":
                self.msges += self.make_icl_demo(x)
            else:
                continue