|
|
from citekit.cite_modules.LLM import LLM |
|
|
from citekit.cite_modules.augment_model import AttributingModule |
|
|
from citekit.pipeline.pipeline import Pipeline, PIPELINE_OUTPUT,PIPELINE_DOC_CACHE |
|
|
from citekit.prompt.prompt import Prompt, ALCEDocPrompt |
|
|
from citekit.Dataset.Dataset import PromptDataset |
|
|
from citekit.evaluator.evaluator import DefaultEvaluator |
|
|
import argparse |
|
|
import json |
|
|
from citekit.utils.utils import cut_and_make_as,one_paragraph,make_as |
|
|
|
|
|
PARA_SEP = '\n\n' |
|
|
if __name__ == '__main__': |
|
|
|
|
|
|
|
|
parser = argparse.ArgumentParser() |
|
|
parser.add_argument("--save_path", type=str, default='res.json', help="Path to the config file") |
|
|
parser.add_argument("--model", type=str, default='gpt-3.5-turbo', help="model name or path") |
|
|
parser.add_argument("--shots", type=int, default=1, help="number of shots") |
|
|
parser.add_argument("--ndoc", type=int, default=3, help="number of docs") |
|
|
parser.add_argument("--pr", action='store_true', help="use cite PR") |
|
|
parser.add_argument("--rouge", action='store_true', help="use rouge") |
|
|
parser.add_argument("--temp", type=float, default=0.5, help="temperature") |
|
|
parser.add_argument("--qa", action='store_true', help="eval qa") |
|
|
parser.add_argument("--mauve", action='store_true', help="eval mauve") |
|
|
parser.add_argument("--length", type=bool, default=True, help="eval length") |
|
|
parser.add_argument("--claims", action='store_true', help="eval claims") |
|
|
parser.add_argument("--qampari", type=str, default=False, help="eval qampari") |
|
|
parser.add_argument("--turns", type=int, default=1, help="k") |
|
|
parser.add_argument("--use_fast_pr", type=str, default=False, help="test") |
|
|
parser.add_argument("--dataset", type=str, default='data/asqa_eval_gtr_top100.json', help="dataset") |
|
|
parser.add_argument("--demo", type=str, default='prompts/AnG.json', help="demo") |
|
|
parser.add_argument("--mode", type=str, default='AnG', help="mode: AnG or plan") |
|
|
args = parser.parse_args() |
|
|
|
|
|
|
|
|
file_path = args.dataset |
|
|
demo_path = args.demo |
|
|
with open(file_path,'r',encoding='utf-8') as file: |
|
|
dataset = json.load(file) |
|
|
with open(demo_path,'r',encoding='utf-8') as file: |
|
|
demo = json.load(file)[args.mode] |
|
|
|
|
|
|
|
|
dataset = PromptDataset(dataset,'question','answer','answers','qa_pairs','claims', docs = lambda data: ALCEDocPrompt().default_load_data_wo_title(data['docs'][:args.ndoc]))[:200] |
|
|
if args.mode == 'AnG': |
|
|
gen_shot = demo['gen_instruction'] + PARA_SEP + demo['gen_shot'] + PARA_SEP |
|
|
answer_ppt = {'INST':demo['gen_instruction'],'shot':gen_shot, 'add':'The next sentence is:'} |
|
|
elif args.mode == 'plan': |
|
|
shot = demo['shot1'] + demo['shot2'] |
|
|
self_ppt = {'INST':demo['INST'],'shot':shot, 'add':'subquestions: \n'} |
|
|
answer_shot = demo['answer_shot_1'] + demo['answer_shot_2'] |
|
|
answer_ppt = {'INST':demo['answer_inst'],'shot':answer_shot,'add':''} |
|
|
|
|
|
prompt = Prompt(template='<shot><INST><question><docs><prefix><sub><span><add>', |
|
|
components={'INST':'{INST}\n\n', |
|
|
'shot':'{shot}', |
|
|
'question':'Question:{question}\n\n', |
|
|
'docs':'{docs}\n', |
|
|
'span':'The highlighted spans are: \n{span}\n\n', |
|
|
'prefix':'Prefix: {prefix}\n\n', |
|
|
'sub':'subquestions: \n{sub}\n\n', |
|
|
'add':'Answer: \n{add}' |
|
|
}) |
|
|
|
|
|
plan_prompt = Prompt(template='<shot><INST><question><docs><sub><add>', |
|
|
components={'INST':'{INST}\n\n', |
|
|
'shot':'{shot}', |
|
|
'question':'Question:{question}\n\n', |
|
|
'docs':'{docs}\n', |
|
|
'sub':'subquestions: \n{sub}\n\n', |
|
|
'add':'{add}'}) |
|
|
|
|
|
|
|
|
evaluator = DefaultEvaluator(args) |
|
|
if args.mode == 'AnG': |
|
|
attribute = AttributingModule(model = args.model) |
|
|
elif args.mode == 'plan': |
|
|
attribute = LLM(model = args.model, prompt_maker = plan_prompt,self_prompt=self_ppt,post_processing=cut_and_make_as('sub')) |
|
|
|
|
|
answer = LLM(model = args.model, prompt_maker = prompt, self_prompt=answer_ppt, share_model_with=attribute.get_first_module(), post_processing=cut_and_make_as('prefix'), iterative=True) |
|
|
if args.mode == 'AnG': |
|
|
attribute.set_target(answer) |
|
|
elif args.mode == 'plan': |
|
|
attribute.set_target(answer,post_processing=cut_and_make_as('sub')) |
|
|
pipeline = Pipeline(save_path=args.save_path, llm = answer, module = attribute, evaluator = evaluator, dataset = dataset) |
|
|
|
|
|
answer.set_output(post_processing=lambda ls: ''.join(map(one_paragraph,ls))) |
|
|
|
|
|
pipeline.run_on_dataset(datakeys=['question','docs'],init_docs='docs',initial_module = attribute) |
|
|
|
