Upload 97 files

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+asset/design.png filter=lfs diff=lfs merge=lfs -text

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 Jiajun Shen
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

app.py

@@ -0,0 +1,104 @@
+from flask import Flask, request, jsonify, send_file, Response
+from flask_cors import CORS
+import openai
+import sys
+import os
+from methods.self_RAG_demo import pipeline, graph
+from citekit.utils.utils import parse_html_config
+
+app = Flask(__name__)
+CORS(app)  # 允许跨域请求
+
+
+@app.route("/")
+def index():
+    return send_file("index.html")
+
+
+@app.route("/run_pipeline", methods=["POST"])
+def run_pipeline():
+    data = request.json  
+    if not data:
+        return jsonify({"error": "Invalid input data"}), 400
+    
+    try:
+        result = pipeline(data)  # 直接调用 pipeline 处理数据
+        print(result)
+        return jsonify(result)  # 返回 JSON 结果
+    except Exception as e:
+        return jsonify({"error": str(e)}), 500
+    
+@app.route("/get_nodes", methods=["POST"])
+def get_nodes(*args, **kwargs):
+    graph.update()
+    try:
+        return jsonify(graph.get_json()) 
+    except Exception as e:
+        return jsonify({"error": str(e)}), 500
+
+@app.route("/update", methods=["POST"])
+def update():
+
+    data = request.json
+    update_info = data.get("update_info")
+    update_object = data.get('update_object')
+    print(update_info, update_object)
+    try:
+        config, update_info = parse_html_config(update_info)
+        print('GOT CONFIG', config, update_info)
+        pipeline.update(update_object, config, update_info)
+        return jsonify({})
+    except Exception as e:
+        return jsonify({"error": str(e)}), 500
+
+
+
+@app.route("/get_config", methods=["POST"])
+def get_config():
+    data = request.json
+    config = data.get("config").lower()
+    module_name = data.get("module_name")
+    module = pipeline.get_module_by_name(module_name)
+
+    try:
+        if config in ['prompt', 'destination', 'max turn', 'global prompt', 'parallel']:
+            return jsonify(module.get_json_config(config))
+        else:
+            raise NotImplementedError
+
+
+    except Exception as e:
+        return jsonify({"error": str(e)}), 500
+
+@app.route("/chat", methods=["POST"])
+def chat():
+    data = request.json
+    api_key = data.get("api_key")
+    user_message = data.get("message")
+
+    if not api_key or not user_message:
+        return jsonify({"error": "API Key and message are required"}), 400
+
+    try:
+        openai.api_key = api_key
+        response = openai.ChatCompletion.create(
+            model="gpt-4o",
+            messages=[
+            {"role": "system", "content": "You are a helpful assistant that follows the instructions of the user. You will be given a pipeline and (maybe) some datapoints in json format. You will be asked questions about the pipeline or the datapoints. Refuse to answer questions that are not about the pipeline or the datapoints."},
+            {"role": "user", "content": user_message}
+            ],
+            stream=True  # 启用流式输出
+        )
+
+        def generate():
+            for chunk in response:
+                if "choices" in chunk and chunk["choices"]:
+                    yield chunk["choices"][0]["delta"].get("content", "")
+
+
+        return Response(generate(), content_type="text/event-stream")  # 使用流式响应
+    except Exception as e:
+        return jsonify({"error": str(e)}), 500
+
+if __name__ == '__main__':
+    app.run(host="0.0.0.0", port=7860)

citekit/Dataset/Dataset.py

@@ -0,0 +1,41 @@
+from torch.utils.data import Dataset
+import json
+
+default_get = lambda key :  lambda data: data[key]
+
+class PromptDataset(Dataset):
+
+    def __init__(self,data_dir,*keys,**projections) -> None:
+        self.data = []
+        for d in data_dir:
+            list_contents = {key:default_get(key)(d) for key in keys if key in d.keys()}
+            dict_contents = {projection:projections[projection](d) for projection in projections.keys()}
+            self.data.append({**list_contents,**dict_contents})
+    
+    def __getitem__(self, index) -> dict:
+    
+        return self.data[index]
+    
+    def __len__(self):
+        return len(self.data)
+
+class FileDataset(PromptDataset):
+
+    def __init__(self,data_dir,*keys,**projections) -> None:
+        with open(data_dir,'r',encoding='utf-8') as file:
+            data_dir = json.load(file)
+        if not keys:
+            keys = data_dir[0].keys()
+
+        self.data = []
+        for d in data_dir:
+            list_contents = {key:default_get(key)(d) for key in keys if key in d.keys()}
+            dict_contents = {projection:projections[projection](d) for projection in projections.keys()}
+            self.data.append({**list_contents,**dict_contents})
+    
+    def __getitem__(self, index) -> dict:
+    
+        return self.data[index]
+    
+    def __len__(self):
+        return len(self.data)

citekit/__init__.py

@@ -0,0 +1 @@
+from __future__ import absolute_import

citekit/attribute/attribute.py

@@ -0,0 +1,222 @@
+import json
+from context_cite import ContextCiter
+import re
+import torch
+from transformers import LlamaForCausalLM, LlamaTokenizer
+
+
+def all_normalize(obj):
+    all_values = []
+    for output_sent_result in obj:
+        for each_doc in output_sent_result:
+            for each_span in each_doc:
+                all_values.append(each_span[1])
+    max_val = max(all_values)
+    min_val = min(all_values)
+    for output_sent_result in obj:
+        for i, each_doc in enumerate(output_sent_result):
+            for j, each_span in enumerate(each_doc):
+                each_span = (each_span[0], (each_span[1] - min_val) / (max_val - min_val))
+                output_sent_result[i][j] = each_span
+    return obj
+
+def all_normalize_in(obj):
+    for output_sent_result in obj:
+        all_values = []
+        for each_doc in output_sent_result:
+            for each_span in each_doc:
+                all_values.append(each_span[1])
+        max_val = max(all_values)
+        min_val = min(all_values)
+        for i, each_doc in enumerate(output_sent_result):
+            for j, each_span in enumerate(each_doc):
+                each_span = (each_span[0], (each_span[1] - min_val) / (max_val - min_val))
+                output_sent_result[i][j] = each_span
+    return obj
+
+def load_json(file_path):
+
+    with open(file_path, 'r') as file:
+        data = file.read()
+    if file_path.endswith('.jsonl'):
+        data = f'[{'},{'.join(data.split("}\n{"))}]'
+    objects = json.loads(data)
+    return objects
+
+def ma(text):
+    pattern = r"Document \[\d+\]\(Title:[^)]+\)"
+
+    match = re.search(pattern, text)
+
+    if match:
+        index = match.end()
+        return index
+    else:
+        return 0
+
+def write_json(file_path, data):
+    with open(file_path, 'w') as json_file:
+        json.dump(data, json_file, indent=4)
+
+
+
+def load_model(model_name_or_path):
+    from transformers import AutoModelForCausalLM, AutoTokenizer
+    model = AutoModelForCausalLM.from_pretrained(
+        model_name_or_path,
+        device_map='auto',
+        token = 'your token'
+    )
+
+    tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)
+    model.eval()
+    return model, tokenizer
+
+
+def compute_log_prob(model, tokenizer, input_text, output_text):
+    inputs = tokenizer(input_text, return_tensors="pt")
+    output_tokens = tokenizer(output_text, return_tensors="pt")["input_ids"]
+    
+    with torch.no_grad():
+        logits = model(**inputs).logits[:, -output_tokens.shape[1]-1:-1, :]
+    
+    log_probs = torch.nn.functional.log_softmax(logits, dim=-1)
+    output_log_probs = log_probs.gather(2, output_tokens.unsqueeze(-1)).squeeze(-1)
+    return output_log_probs.sum().item()
+
+def compute_contributions(model, tokenizer, question, docs, output):
+    full_input = question + '\n\n' + '\n'.join(docs)
+    base_prob = compute_log_prob(model, tokenizer, full_input, output)
+    
+    contributions = []
+    for i in range(len(docs)):
+        reduced_docs = docs[:i] + docs[i+1:]
+        reduced_input = question + '\n\n' + '\n'.join(reduced_docs)
+        reduced_prob = compute_log_prob(model, tokenizer, reduced_input, output)
+        contributions.append(base_prob - reduced_prob)
+    
+    return contributions
+    
+class InterpretableAttributer:
+
+    def __init__(self, levels=['doc', 'span', 'word'], model = 'gpt-2'):
+        for level in levels:
+            assert level in ['doc', 'span', 'word'], f'Invalid level: {level}'
+        # span before doc
+        self.levels = sorted(levels, key=lambda x: ['span', 'doc', 'word'].index(x))
+        #self.model, self.tokenizer = load_model(model)
+
+
+    def attribute(self, question, docs, output):
+        attribute_results = {}
+        for level in self.levels:
+            attribute_result = []
+            for sentence in output:
+                attribute_result.append(self._attribute(question, docs, sentence, level))
+            attribute_results[level] = attribute_result
+        return attribute_results
+    
+
+    def _attribute(self, question, docs, output, level):
+        if level == 'doc':
+            return self.doc_level_attribution(question, docs, output)
+        elif level == 'span':
+            return self.span_level_attribution(question, docs, output)
+        elif level == 'word':
+            return self.word_level_attribution(question, docs, output)
+        else:
+            raise ValueError(f'Invalid level: {level}')
+    
+    def span_level_attribution(self, question, docs, output):
+        # USE CONTEXT CITE
+        context = '\n\n'.join(docs)
+        response = output
+
+        cc = ContextCiter(self.model, self.tokenizer, context, question)
+        _, prompt = cc._get_prompt_ids(return_prompt=True)
+        cc._cache["output"] = prompt + response  
+        result = cc.get_attributions(as_dataframe=True, top_k=1000).data.to_dict(orient='records')
+        return result
+    
+    
+    def parse_attribution_results(self, docs, results):
+        context = '\n\n'.join(docs)
+        lens = [len(doc) for doc in docs]
+        len_sep = len('\n\n')
+        final_results = {}
+        for level, result in results.items(): 
+            if level == 'span':
+                ordered_all_sents = []
+                for output_sent_result in result:
+                    final_end_for_span = {}
+                    all_span_results = []
+                    for each_span in output_sent_result:
+                        span_text = each_span["Source"]
+                        span_score = each_span["Score"]
+                        start = 0
+                        if span_text in final_end_for_span:
+                            start = final_end_for_span[span_text]
+                        span_start = context.find(span_text, start)
+                        span_end = span_start + len(span_text)
+                        final_end_for_span[span_text] = span_end
+                        # locate the document
+                        doc_idx = 0
+                        while span_start > lens[doc_idx]:
+                            span_start -= lens[doc_idx] + len_sep
+                            span_end -= lens[doc_idx] + len_sep
+                            doc_idx += 1
+                        all_span_results.append((span_start, span_score, doc_idx))
+                    ordered = [[] for _ in range(len(docs))]
+                    for span_start, span_score, doc_idx in all_span_results:
+                        ordered[doc_idx].append((span_start, span_score))
+                    for i in range(len(docs)):
+                        doc = docs[i]
+                        real_start = ma(doc)
+                        ordered[i] = sorted(ordered[i], key=lambda x: x[0])
+                        ordered[i][0] = (real_start, ordered[i][0][1])
+
+                    ordered_all_sents.append(ordered)
+                final_results[level+'_level'] = all_normalize_in(ordered_all_sents)
+            elif level == 'doc':
+                self.span_to_doc(result)
+            else:
+                raise NotImplementedError(f'Parsing for {level} not implemented yet')
+        return final_results
+    
+    def span_to_doc(self, results):
+        import numpy as np
+        span_level = results['span_level']
+        doc_level = []
+        for output_sent_result in span_level:
+            doc_level.append([np.mean([span[1] for span in doc]) for doc in output_sent_result])
+        results['doc_level'] = doc_level
+        
+
+    def attribute_for_result(self, result):
+        docs = result['doc_cache']
+        question = result['data']['question']
+        output = result['output']
+        attribution_results = self.attribute(question, docs, output)
+        parsed_results = self.parse_attribution_results(docs, attribution_results)
+        result.update(parsed_results)
+
+        if 'doc' not in self.levels:
+            # if doc is not in the levels, we need to convert the span level to doc level
+            print('Converting span level to doc level...')
+            try:
+                self.span_to_doc(result)
+                print('Conversion successful')
+            except Exception as e:
+                print(f'Error converting span level to doc level: {e}')
+
+    def attribute_for_results(self, results):
+        for result in results:
+            self.attribute_for_result(result)
+        return results
+    
+
+if __name__ == '__main__':
+    attributer = InterpretableAttributer(levels=['span'])
+    results = load_json('res_attr.json')
+    attributer.attribute_for_results(results)
+    write_json('res_attr_span.json', results)

citekit/cite_modules/LLM.py

@@ -0,0 +1,427 @@
+import torch
+from citekit.prompt.prompt import Prompt
+import re
+from citekit.utils.utils import one_paragraph, first_sentence, make_as
+import random
+import os
+
+
+
+class Module:
+    module_count = 1
+    def __init__(self,prompt_maker: Prompt = None, pipeline = None, self_prompt = {}, iterative = False, merge = False, max_turn =6, output_as = None, parallel = False) -> None:
+        self.self_prompt = self_prompt
+        self.use_head_prompt = True
+        self.connect_to(pipeline)
+        self.prompt_maker = prompt_maker
+        self.last_message = ''
+        self.destinations = []
+        self.conditions = {}
+        self.head_key = None
+        self.parallel = parallel
+        self.iterative = iterative
+        self.merge = merge
+        self.head_process = one_paragraph
+        self.max_turn = max_turn
+        self.multi_process = False
+        self.output_cond = {} # {cond : {'post_processing':post, 'end':end}}
+        self.count = Module.module_count
+        Module.module_count += 1
+        self.if_add_output_to_head = False
+        self.turns = 0
+        self.end = False
+
+    def __str__(self) -> str:
+        if self.model_type:
+            return f'{self.model_type}-[{self.count}]'
+        else:
+            return f'Unknown-type module-[{self.count}]'
+        
+    def get_json_config(self, config):
+        print('get_json_config:',config)
+        avaliable_mapping = {
+            'max turn': 'max_turn',
+            'prompt': 'prompt',
+            'destination': 'destination',
+            'global prompt': 'head_key',
+        }
+        if config == 'prompt':
+            prompt_info = {
+                'template': self.prompt_maker.template,
+                'components': self.prompt_maker.components
+            }
+            self_info = self.self_prompt
+            
+            return {
+                'prompt_info': prompt_info,
+                'self_info': self_info
+            }
+        elif config == 'destination':
+            return {
+                'destination': str(self.destinations[0])
+            }
+        elif config in ['max turn','global prompt']:
+            config = avaliable_mapping[config]
+            print('getting the config:',config)
+            return getattr(self, config)
+        else:
+            raise NotImplementedError(f'get_json_config for {config} is not implemented')
+    
+    def get_destinations(self):
+        return self.destinations
+    
+    def update(self, config, update_info):
+
+        if config == 'prompt':
+            template = update_info['template']
+            components = update_info['components']
+            self_prompt = update_info['self_prompt']
+            import copy
+            # avoid changing the original prompt_maker
+            self.prompt_maker = copy.deepcopy(self.prompt_maker)
+
+            self.prompt_maker.update(template=template, components=components)
+            self.self_prompt = self_prompt
+
+        elif config == 'destination':
+            print('update destination:',update_info[0], 'post_processing:',update_info[1])
+            if update_info[1] == 'None':
+                self.set_target(update_info[0])
+            else:
+                self.set_target(update_info[0], post_processing=make_as(update_info[1]))
+
+        elif config == 'delete_destination':
+            for i, d in enumerate(self.destinations):
+                if str(d) == str(update_info):
+                    self.destinations.remove(d)
+                    del self.conditions[d]
+                    break
+        elif config == 'header':
+            self.add_to_head(update_info, sub = True)
+        elif config == 'max turn':
+            self.max_turn = update_info
+        else:
+            raise NotImplementedError(f'update for {config} is not implemented')
+
+    def end_multi(self):
+        return
+
+    def set_use_head_prompt(self,use):
+        assert isinstance(use,bool)
+        self.use_head_prompt = use
+    
+    def reset(self):
+        self.end = False
+        self.turns = 0
+    
+    def change_to_multi_process(self,bool_value):
+        if bool_value:
+            self.last_message = []
+        else:
+            self.last_message = ''
+        self.multi_process = bool_value
+    @property
+    def get_use_head_prompt(self):
+        return self.use_head_prompt
+
+    def generate(self, head_prompt: dict = {}, dynamic_prompt: dict = {}):
+        raise NotImplementedError
+    
+    def send(self):
+        for destination in self.destinations:
+            cond = self.conditions[destination]['condition']
+            if cond(self):
+                return destination
+        return None
+
+    def set_target(self,destination, condition = lambda self: True, post_processing = lambda x:x) -> None:
+        self.conditions[destination] = {'condition': condition, 'post_processing' : post_processing}
+        self.destinations = [destination] + self.destinations
+        destination.connect_to(self.pipeline)
+    
+    def clear_destination(self):
+        self.destinations = []
+        self.conditions = {}
+
+    def add_output_to_head(self, outputs):
+        if self.if_add_output_to_head:
+            if not self.head_sub:
+                if self.head_key not in self.pipeline.head.keys():
+                    self.pipeline.head.update({self.head_key: self.head_process(outputs)})
+                else:
+                    self.pipeline.head[self.head_key] += '\n'
+                    self.pipeline.head[self.head_key] += self.head_process(outputs)
+            else: 
+                self.pipeline.head[self.head_key] = self.head_process(outputs)
+
+    def connect_to(self, pipeline = None) -> None:
+        self.pipeline = pipeline
+        if pipeline:
+            pipeline.module.append(self)
+
+    def output(self):
+        outed = False
+        for cond, post_and_end in self.output_cond.items():
+            if cond(self):
+                if not outed:
+                    if not self.merge:
+                        self.pipeline.output.append(post_and_end['post_processing'](self.last_message))
+                    else:
+                        self.pipeline.output.append(post_and_end['post_processing'](''.join(self.last_message)))
+                    outed = True
+                if post_and_end['end']:
+                    self.end =  True
+
+    def set_output(self, cond = lambda self: True, post_processing = lambda x:x, end = True):
+        self.output_cond[cond] = {'post_processing': post_processing, 'end' : end}
+    
+    def get_first_module(self):
+        return self
+    
+    def add_to_head(self, datakey, sub = False, process = None):
+        self.if_add_output_to_head = True
+        self.head_key = datakey
+        self.head_sub = sub
+        if process:
+            self.head_process = process
+
+
+def load_model(model_name_or_path,dtype = torch.float16):
+    from transformers import AutoModelForCausalLM, AutoTokenizer
+    model = AutoModelForCausalLM.from_pretrained(
+        model_name_or_path,
+        torch_dtype=dtype,
+        device_map='auto',
+    )
+
+    tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)
+    model.eval()
+    return model, tokenizer
+
+
+class LLM(Module):
+    model_type = 'Generator'
+    def __init__(self, model = None, prompt_maker: Prompt =None, pipeline = None, post_processing = None, self_prompt = {}, device = 'cpu',temperature = 0.5 ,stop = None, max_turn = 6, share_model_with = None, iterative = False, auto_cite = False, output = None,merge = False, noisy = True, parallel = False, output_as ='Answer', auto_cite_from = 'docs') -> None:
+        super().__init__(prompt_maker,pipeline,self_prompt, iterative, merge, parallel = parallel)
+        self.max_turn = max_turn
+        if post_processing:
+            self.post_processing = post_processing
+        else:
+            self.post_processing = lambda x: {output_as:x}
+        if model:
+            self.model_name = model
+        self.stop = stop
+        self.multi_process = False
+        self.noisy = noisy
+        self.head_process = one_paragraph
+        self.auto_cite = auto_cite
+        if auto_cite:
+            self.cite_from = auto_cite_from
+        if model:
+            if 'gpt' not in model.lower():
+                if not share_model_with:
+                    print('loading model...')
+                    self.model, self.tokenizer = self.load_model(model)
+                else: 
+                    print('sharing model...')
+                    self.model, self.tokenizer = share_model_with.model, share_model_with.tokenizer
+                self.temperature = temperature
+                self.device = device
+            else:
+                self.openai_key = os.getenv('OPENAI_API_KEY')
+        self.output_cond = {} # {cond : {'post_processing':post, 'end':end}}
+        self.if_add_output_to_head = False
+
+        self.token_used = 0
+
+    def reset(self):
+        self.end = False
+        self.turns = 0
+        self.token_used = 0
+        
+    
+    def __str__(self) -> str:
+        if self.model_name:
+            return f'{self.model_name}-[{self.count}]'
+        else:
+            return 'unknown model'
+    
+    def __repr__(self) -> str:
+        return (f'{self.prompt_maker}\n|\n|\nV\n{self}\n|\n|\nV\n'+ '/'.join([str(des) for des in self.destinations]+['output']))
+    
+    def load_model(self, model_name_or_path,dtype = torch.float16):
+        from transformers import AutoModelForCausalLM, AutoTokenizer
+        model = AutoModelForCausalLM.from_pretrained(
+        model_name_or_path,
+        torch_dtype=dtype,
+        device_map='auto',
+    )
+
+        tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)
+        model.eval()
+        return model, tokenizer
+
+    def set_cite(self,key):
+        self.cite_from = key
+        self.auto_cite = True
+    
+    def generate_content(self, prompt):
+        if 'gpt' in self.model_name.lower():
+            import openai
+            openai.api_key = self.openai_key
+            prompt = [
+                    {'role': 'system',
+                     'content': "You are a good helper who follow the instructions"},
+                    {'role': 'user', 'content': prompt}
+                ]
+            response = openai.ChatCompletion.create(
+                model=self.model_name,
+                messages=prompt,
+                max_tokens=500,
+                stop = self.stop
+            )
+            self.token_used += response['usage']['completion_tokens'] + response['usage']['prompt_tokens']
+            return response['choices'][0]['message']['content']
+        
+        else:
+            inputs = self.tokenizer(prompt, return_tensors="pt").to(self.model.device)
+            stop = [] if self.stop is None else self.stop
+
+            outputs = self.model.generate(
+                    **inputs,
+                    do_sample = True,
+                    max_new_tokens = 200,
+                    temperature = self.temperature
+                    )
+            self.token_used += len(outputs[0])
+
+            outputs = self.tokenizer.decode(outputs[0][inputs['input_ids'].size(1):], skip_special_tokens=True)
+            return one_paragraph(outputs)
+            print(outputs)
+
+
+    def generate(self, head_prompt: dict = {}, dynamic_prompt: dict = {}):
+        if self.use_head_prompt: 
+            #print(head_prompt,self.self_prompt,dynamic_prompt)
+            prompt = self.prompt_maker(head_prompt,self.self_prompt,dynamic_prompt)
+        else:
+            prompt = self.prompt_maker(self.self_prompt,dynamic_prompt)
+        if self.noisy:
+            print(f'prompt to {str(self)}:\n',prompt,'\n\n')
+        self.turns += 1
+        
+        outputs = self.generate_content(prompt)
+        #print('DEBUG:',outputs)
+        if self.noisy:
+            print('OUTPUT:')
+            print(outputs)
+        if self.auto_cite:
+            outputs = self.cite_from_prompt({**head_prompt,**self.self_prompt,**dynamic_prompt},outputs)
+        if self.multi_process:
+            self.last_message.append(outputs)
+        else:
+            self.last_message = outputs
+
+
+        self.add_output_to_head(outputs)
+
+        destination = self.send()
+
+        if self.turns > self.max_turn:
+            self.end = True
+        if destination in self.conditions:
+            return self.conditions[destination]['post_processing'](outputs)
+        else: 
+            return self.post_processing(outputs)
+        
+    def add_output_to_head(self, outputs):
+        if self.if_add_output_to_head:
+            if not self.head_sub:
+                if self.head_key not in self.pipeline.head.keys():
+                    self.pipeline.head.update({self.head_key: self.head_process(outputs)})
+                else:
+                    self.pipeline.head[self.head_key] += '\n'
+                    self.pipeline.head[self.head_key] += self.head_process(outputs)
+            else: 
+                self.pipeline.head[self.head_key] = self.head_process(outputs)
+
+    def output(self):
+        outed = False
+        for cond, post_and_end in self.output_cond.items():
+            if cond(self):
+                if not outed:
+                    if not self.merge and not self.iterative:
+                        self.pipeline.output.append(post_and_end['post_processing'](self.last_message))
+                    else:
+                        self.pipeline.output.append(post_and_end['post_processing'](' '.join(self.last_message)))
+                    outed = True
+                if post_and_end['end']:
+                    self.end =  True
+
+    def set_output(self, cond = lambda self: True, post_processing = lambda x:x, end = True):
+        self.output_cond[cond] = {'post_processing': post_processing, 'end' : end}
+        
+
+    def cite_from_prompt(self,prompt_dict,input):
+        input = first_sentence(input)
+        cite_docs = prompt_dict[self.cite_from]
+        refs = re.findall(r'\[\d+\]', cite_docs)
+        pattern = r'([.!?])\s*$'
+        if refs:
+            cite = ''.join(refs)
+        else:
+            cite = ''
+        output = re.sub(pattern, rf' {cite}\1 ', input)
+        if cite not in output:
+            output += cite
+        return output
+    def add_to_head(self, datakey, sub = False, process = None):
+        self.if_add_output_to_head = True
+        self.head_key = datakey
+        self.head_sub = sub
+        if process:
+            self.head_process = process
+
+
+
+class TestLLM(LLM):
+    def __init__(self, model='gpt-4', prompt_maker: Prompt = None, pipeline=None, post_processing=lambda x: x, self_prompt={}, device='cpu', temperature=0.5, stop=None, max_turn=6,share_model_with = None, iterative= False, ans = None) -> None:
+        super().__init__(model,prompt_maker,pipeline,self_prompt=self_prompt,share_model_with=share_model_with,iterative=iterative)
+        self.max_turn = max_turn
+        self.post_processing = post_processing
+        self.model_name = model
+        self.last_message = ''
+        self.stop = stop
+        self.output_cond = {} # {cond : {'post_processing':post, 'end':end}}
+        self.if_add_output_to_head = False
+
+        self.token_used = 0
+        self.ans = 'Strain[1], turns:, heat[2][4]. Sent2[5]. Sent3.\n\n rdd' if not ans else ans
+    def generate_content(self, prompt):
+        return self.ans
+
+
+class AutoAISLLM(LLM):
+    def __init__(self, model=None, prompt_maker: Prompt = None, pipeline=None, post_processing=None, self_prompt={}, device='cpu', temperature=0.5, stop=None, max_turn=6, share_model_with=None, iterative=False, auto_cite=False, output=None, merge=False, noisy=False, output_as='Answer') -> None:
+        super().__init__(model, prompt_maker, pipeline, post_processing, self_prompt, device, temperature, stop, max_turn, share_model_with, iterative, auto_cite, output, merge, noisy, output_as)
+
+        self.prompt_maker = Prompt('<INST><premise><claim>\n Answer: ',components={
+            'INST':'{INST}\n\n',
+            'premise':'Premise: {premise}\n\n',
+            'claim':'Claim: {claim}\n',
+        })
+        self.self_prompt={'INST': 'In this task, you will be presented a premise and a claim. If the premise entails the claim, output "1", otherwise output "1". Your answer should only contains one number without any other letters and punctuations.'}
+    
+    def generate(self, premise, claim):
+        dict_answer = super().generate({'premise':premise,'claim':claim})
+        return dict_answer.get('Answer')
+
+
+
+if __name__ == '__main__':
+    prompt = Prompt(template='<INST><Question><Docs><feedback><Answer>',components={'INST':'{INST}\n\n', 
+                                                                                    'Question':'Question:{Question}\n\n',
+                                                                                    'Docs':'{Docs}\n',
+                                                                                    'feedback':'Here is the feed back of your last response:{feedback}\n',
+                                                                                    'Answer':'Here is answer and you have to give feedback:{Answer}'})
+    m = LLM('gpt')

citekit/cite_modules/Retrieve.py

@@ -0,0 +1,99 @@
+import argparse
+import csv
+import json
+import os
+import time
+import pickle
+
+import numpy as np
+import torch
+from tqdm import tqdm
+from rank_bm25 import BM25Okapi
+from sentence_transformers import SentenceTransformer
+
+def gtr_build_index(encoder, docs):
+    with torch.inference_mode():
+        embs = encoder.encode(docs, show_progress_bar=True, normalize_embeddings=True)
+        embs = embs.astype("float16")
+
+    GTR_EMB = os.environ.get("GTR_EMB")
+    with open(GTR_EMB, "wb") as f:
+        pickle.dump(embs, f)
+    return embs
+
+
+class DPRRetriever:
+    def __init__(self, DPR_WIKI_TSV, GTR_EMB = None, emb_model = "sentence-transformers/gtr-t5-xxl") -> None:
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.device = device
+        self.encoder = SentenceTransformer(emb_model, device = device)
+        self.docs = []
+        print("loading wikipedia file...")
+        with open(DPR_WIKI_TSV) as f:
+            reader = csv.reader(f, delimiter="\t")
+            for i, row in enumerate(reader):
+                if i == 0:
+                    continue
+                self.docs.append(row[2] + "\n" + row[1])
+        if not GTR_EMB:
+            print("gtr embeddings not found, building...")
+            embs = gtr_build_index(self.encoder, self.docs)
+        else:
+            print("gtr embeddings found, loading...")
+            with open(GTR_EMB, "rb") as f:
+                embs = pickle.load(f)
+
+        self.gtr_emb = torch.tensor(embs, dtype=torch.float16, device=device)
+
+    def retrieve(self, question, topk):
+        with torch.inference_mode():
+            query = self.encoder.encode(question, batch_size=4, show_progress_bar=True, normalize_embeddings=True)
+            query = torch.tensor(query, dtype=torch.float16, device=self.device)
+        query = query.to(self.device)
+        scores = torch.matmul(self.gtr_emb, query)
+        score, idx = torch.topk(scores, topk)
+        ret = []
+        for i in range(idx.size(0)):
+            title, text = self.docs[idx[i].item()].split("\n")
+            ret.append({"id": str(idx[i].item() + 1), "title": title, "text": text, "score": score[i].item()})
+        
+        return ret
+    
+    def __repr__(self) -> str:
+        return 'DPR Retriever'
+    
+    def __str__(self) -> str:
+        return repr(self)
+
+class BM25Retriever:
+    def __init__(self, DPR_WIKI_TSV):
+        self.docs = []
+        self.tokenized_docs = []
+        print("loading wikipedia file...")
+        with open(DPR_WIKI_TSV) as f:
+            reader = csv.reader(f, delimiter="\t")
+            for i, row in enumerate(reader):
+                if i == 0:
+                    continue
+                self.docs.append(row[2] + "\n" + row[1])
+                self.tokenized_docs.append((row[2] + " " + row[1]).split())
+
+        print("BM25 index building...")
+        self.bm25 = BM25Okapi(self.tokenized_docs)
+
+    def retrieve(self, question, topk):
+        query = question.split()
+        scores = self.bm25.get_scores(query)
+        topk_indices = scores.argsort()[-topk:][::-1]
+        ret = []
+        for idx in topk_indices:
+            title, text = self.docs[idx].split("\n", 1)
+            ret.append({"id": str(idx + 1), "title": title, "text": text, "score": scores[idx]})
+        
+        return ret
+    def __repr__(self) -> str:
+        return 'BM25 Retriever'
+    
+    def __str__(self) -> str:
+        return repr(self)
+    

citekit/cite_modules/augment_model.py

@@ -0,0 +1,455 @@
+from citekit.cite_modules.LLM import Module,LLM
+from citekit.cite_modules.Retrieve import DPRRetriever
+from citekit.evaluator.evaluator import _run_nli_autoais, Evaluator
+from citekit.prompt.prompt import Prompt
+from citekit.utils.utils import one_paragraph, make_as
+from sentence_transformers import SentenceTransformer
+import re
+import random
+
+
+class Retriever(Module):
+    model_type = 'retriever'
+    def __init__(self, documents = None ,retrieve_by = 'index', prompt_maker = None, pipeline = None, post_processing = lambda input, output: {'RetrievedDocs':output}, self_prompt = {},topk = 3,add_id = True, merge = False, tsv_path = 'None', emb_path = 'None', retrieve_from_data = True, parallel = False) -> None:
+        super().__init__(prompt_maker,pipeline,self_prompt,merge=merge, parallel=parallel)
+        self.retrieve_by = retrieve_by
+        self.use_head_prompt = False
+        self.dataset_documents = None
+        self.documents = None
+        self.default_doc_key = 'docs'
+        self.retrieve_from_data = retrieve_from_data
+        if not documents:
+            self.documents = self.pipeline.doc_cache
+        else:
+            self.dataset_documents = documents
+        self.post_processing = post_processing
+        self.if_add_output_to_head = False
+        self.topk = topk
+        self.add_id = add_id
+        if retrieve_by =='bm25':
+            self.bm25_module_loaded = False
+            from rank_bm25 import BM25Okapi
+            import nltk
+            nltk.download('punkt')
+            from nltk.tokenize import word_tokenize
+            self.word_tokenize = word_tokenize
+            self.BM25Okapi = BM25Okapi
+            self.bm25_module_loaded = True
+        elif retrieve_by == 'dpr':
+            self.dpr_retriever = DPRRetriever(DPR_WIKI_TSV=tsv_path, 
+                        GTR_EMB=emb_path)
+    
+    def generate(self,head_prompt: dict = {}, dynamic_prompt: dict = {}):
+        index = self.pipeline.data_index
+        if self.dataset_documents:
+            if isinstance(self.dataset_documents[0], list):
+                # Each data has a document
+                self.documents = self.dataset_documents[index]
+            else:
+                self.documents = self.dataset_documents
+        else:
+            if not self.retrieve_from_data:
+                self.documents = self.pipeline.doc_cache.show_docs()
+            else:
+                def _stringtfy(doc):
+                    if isinstance(doc, str):
+                        return doc
+                    return f"({doc['title']}) {doc['text']}"
+                self.documents = list(map(_stringtfy, self.pipeline.current_data[self.default_doc_key]))
+        # query
+        if self.use_head_prompt: 
+            prompt = self.prompt_maker(head_prompt,self.self_prompt,dynamic_prompt)
+        else:
+            prompt = self.prompt_maker(self.self_prompt,dynamic_prompt)   
+
+        retrieved_docs = []
+        if self.retrieve_by == 'index':
+            # query : Document [1][2]
+            indice = [int(r[1:]) - 1 for r in re.findall(r"\[\d+",prompt)]
+            for index in indice:
+                retrieved_docs.append(self.documents[index])
+            if len(retrieved_docs) > self.topk:
+                retrieved_docs = retrieved_docs[:self.topk]
+        elif self.retrieve_by =='bm25':
+            # natural language query
+            tokenized_docs = [self.word_tokenize(doc.lower()) for doc in self.documents]
+            bm25 = self.BM25Okapi(tokenized_docs)
+            tokenized_query = self.word_tokenize(prompt.lower()) 
+            doc_scores = bm25.get_scores(tokenized_query)
+            if self.topk > len(doc_scores):
+                self.topk = len(doc_scores) - 1
+            top_k_idx = sorted(range(len(doc_scores)), key=lambda i: doc_scores[i], reverse=True)[:self.topk]
+            retrieved_docs = [self.documents[index] for index in top_k_idx]
+            retrieved_docs_new = []
+            for re_doc in retrieved_docs:
+                self.pipeline.doc_cache.add_doc(re_doc,self.add_id)
+                retrieved_docs_new.append(self.pipeline.doc_cache.get_last())
+            retrieved_docs = retrieved_docs_new
+            #raise NotImplementedError
+        
+
+        elif self.retrieve_by =='gtr':
+            docs_dict  = self.dpr_retriever.retrieve(prompt,topk=self.topk)
+            retrieved_docs = [f"({d['title']}) {d['text']}" for d in docs_dict]
+            retrieved_docs_new = []
+            for re_doc in retrieved_docs:
+                self.pipeline.doc_cache.add_doc(re_doc,self.add_id)
+                retrieved_docs_new.append(self.pipeline.doc_cache.get_last())
+            retrieved_docs = retrieved_docs_new
+        
+        retrieved_docs_prompt = '\n'.join(retrieved_docs)
+        destination = self.send()
+        if self.multi_process:
+            self.last_message.append(retrieved_docs_prompt)
+        else:
+            self.last_message = retrieved_docs_prompt
+        #print(self.last_message)
+
+        if self.if_add_output_to_head:
+            self.pipeline.head.update({self.head_key:retrieved_docs_prompt})
+        if destination in self.conditions:
+            try:
+                return self.conditions[destination]['post_processing'](prompt,retrieved_docs_prompt)
+            except:
+                return self.conditions[destination]['post_processing'](retrieved_docs_prompt)
+            
+        else:
+            return retrieved_docs_prompt
+        raise NotImplementedError
+
+
+class EvalModule(Module, Evaluator):
+    model_type = 'evaluator'
+    def __init__(self, prompt_maker = None, pipeline=None, self_prompt={},criteria = None, iterative = False, max_turn =6 ,parallel = False) -> None:
+        Module.__init__(self,prompt_maker, pipeline, self_prompt,iterative=iterative,max_turn=max_turn, parallel=parallel)
+        Evaluator.__init__(self,criteria,pipeline)
+    
+    def generate(self, head_prompt: dict = {}, dynamic_prompt: dict = {}):
+        result = {}
+        p_data = {**head_prompt, **self.self_prompt,**dynamic_prompt}
+        for criteria, get_data in self.get_data.items():
+            data_dict = {}
+            for k, v in get_data.items():
+                if v == 'doc_cache':
+                    data_dict[k] = self.pipeline.doc_cache.show_docs()
+                else:
+                    if v in p_data.keys():
+                        data_dict[k] = p_data[v]
+                    else:
+                        data_dict[k] = self.pipeline.current_data[v]
+            eval_func = Evaluator.eval_criteria[criteria]
+            data = [data_dict]
+            result[criteria] = eval_func(data)
+
+        if self.multi_process:
+            self.last_message.append(result)
+        else:
+            self.last_message = result
+        destination = self.send()
+        if destination in self.conditions:
+            return self.conditions[destination]['post_processing'](result)
+        else:
+            return result
+
+
+class CitationSimplyfier(Module):
+    '''
+    Simplify the citation of the 'answer' in prompt.  
+    Argument can be changed to fit into different name of key in prompts 
+    By Defaut, the simplifier simplifies the 'answer' and output the sring with citation simplified.
+    '''
+    model_type = 'simplifier' 
+    def __init__(self, prompt_maker = None, pipeline=None, self_prompt={}, criteria = None, key = 'answer', test = False, allow_empty = True) -> None:
+        Module.__init__(self,prompt_maker, pipeline, self_prompt)
+        if not test: 
+            self.entail = _run_nli_autoais
+        else:
+            self.entail = lambda p,c : random.randint(0,1)
+            self.docs = ['0'] * 100
+        self.key = key
+        self.allow_empty = allow_empty
+    def generate(self, head_prompt: dict = {}, dynamic_prompt: dict = {}) -> str:
+        docs = self.pipeline.doc_cache
+        prompt = {**head_prompt, **dynamic_prompt}
+        answer = prompt[self.key]
+
+        refs = re.findall(r'\[\d+\]', answer)
+        last_ref_index = None
+        for match in re.finditer(r'\[\d+\]', answer):
+            last_ref_index = match.end()
+    
+        if not refs:
+            return answer
+    
+        refs_str = ''.join(refs)
+
+        def simplify(sentence, refs, docs):
+            ref_numbers = [int(ref.strip('[]')) for ref in refs]
+            
+            docs_combined = ''.join(docs[ref - 1] for ref in ref_numbers if ref - 1 < len(docs))
+            
+            if not self.entail(docs_combined, sentence):
+                if self.allow_empty:
+                    return ''
+                return ''.join(refs)  
+            
+            if len(ref_numbers) == 1:
+                return ''.join(f'[{num}]' for num in ref_numbers)
+            
+            def remove_and_test(ref_numbers):
+                for i, ref in enumerate(ref_numbers):
+                    new_ref_numbers = ref_numbers[:i] + ref_numbers[i+1:]
+                    new_docs_combined = ''.join(docs[r - 1] for r in new_ref_numbers if r - 1 < len(docs))
+                    if self.entail(new_docs_combined, sentence):
+                        if len(new_ref_numbers) == 1:
+                            return new_ref_numbers
+                        return remove_and_test(new_ref_numbers)
+                return ref_numbers
+
+            simplified_ref_numbers = remove_and_test(ref_numbers)
+            
+            simplified_refs = ''.join(f'[{num}]' for num in simplified_ref_numbers)
+            return simplified_refs
+
+        simplified_refs = simplify(answer,refs,docs)
+
+        sentence_without_refs = re.sub(r'\[\d+\]', '', answer)
+        
+        if last_ref_index is not None:
+            output = sentence_without_refs[:last_ref_index - len(refs_str)] + simplified_refs + sentence_without_refs[last_ref_index - len(refs_str):]
+        else:
+            output =  sentence_without_refs + simplified_refs
+
+        if not simplified_refs and self.allow_empty:
+            output = ''
+
+            
+        if self.multi_process:
+            self.last_message.append(output)
+        else:
+            self.last_message = output
+        
+        return output
+
+
+class Verifier(Module):
+
+    '''
+    Verifier is currently only used for single sentence/single target answer, not for parallel or iterative answer.
+    Verifier will return dynamic prompt, not like other modules returning output. It is a judger only to decide the target module.
+    By default, the verifoer verifies whether the documents support the answer.
+    '''
+    model_type = 'verifier'
+    def __init__(self, prompt_maker = None, pipeline=None, self_prompt={}, criteria = None, key = 'answer', test = False) -> None:
+        Module.__init__(self,prompt_maker, pipeline, self_prompt)
+        if not test: 
+            self.entail = _run_nli_autoais
+        else:
+            self.entail = lambda p,c : random.randint(0,1)
+            self.docs = ['s'] * 100
+        self.key = key
+        self.test = test
+
+    # Overcite this function to 
+    def verifier_judge(self,**kargs):
+        docs = self.pipeline.doc_cache
+        answer = kargs[self.key]
+        refs = re.findall(r'\[\d+\]', answer)
+        if not refs:
+            return False
+        ref_numbers = [int(ref.strip('[]')) for ref in refs]
+        
+        docs_combined = ''.join(docs[ref - 1] for ref in ref_numbers if ref - 1 < len(docs))
+        return bool(self.entail(docs_combined, re.sub(r'\[\d+\]', '', answer)))
+
+
+    def generate(self, head_prompt: dict = {}, dynamic_prompt: dict = {}):
+        prompt = {**head_prompt, **dynamic_prompt}
+        out = self.verifier_judge(**prompt)
+
+        self.last_message = out
+
+        self.turns += 1
+        return dynamic_prompt
+
+
+class AugmentCluster():
+    def __init__(self, module_list = []) -> None:
+        self.module_list = module_list
+        module_count = len(module_list)
+        for i in range(module_count - 1):
+            assert isinstance(module_list[i],LLM) and isinstance(module_list[i+1],LLM)
+            module_list[i].set_target(module_list[i+1], post_processing = module_list[i].post_processing)
+    
+    def get_first_module(self):
+        return self.module_list[0]
+    
+    def get_destinations(self):
+        return self.module_list[-1].get_destinations()
+    
+    def reset(self):
+        for module in self.module_list:
+            module.reset()
+
+    def update(self, config, update_info):
+        print('updating the AugmentCluster', update_info)
+        self.module_list[-1].update(config, update_info)
+
+    def __str__(self):
+        print('getting str of the cluster', ' -> '.join([str(module) for module in self.module_list]))
+        return ' -> '.join([str(module) for module in self.module_list])
+    
+    def set_target(self,destination, condition = lambda self: True, post_processing = lambda x: x) -> None:
+        self.module_list[-1].set_target(destination, condition, post_processing)
+    
+    def set_output(self, cond = lambda self: True, post_processing = lambda x:x, end = True):
+        self.module_list[-1].set_output(cond, post_processing, end)
+
+    def connect_to(self, pipeline = None) -> None:
+        for module in self.module_list:
+            module.connect_to(pipeline)
+        pipeline.stored_clusters.append(self)
+
+
+class Attribute_post_select(LLM):
+    noisy = False
+    model_name = 'function'
+    def generate(self, head_prompt: dict = {}, dynamic_prompt: dict = {}):
+        print('post_select', head_prompt, dynamic_prompt)
+        docs = head_prompt['docs']
+        ans_docs = one_paragraph(dynamic_prompt['span']).split("\n")
+        spans = [ans_doc[14:].split("<SPAN_DELIM>") for ans_doc in ans_docs]
+        msg = ''
+        span_list = {}
+        doc_map = {}
+        j = 1
+        i = 1
+        for doc in spans:
+            if doc!= [] :
+                span_list[f'{i}'] = []
+                msg += f'Document [{i}]:\n'
+                for span in doc:
+                    if len(span)> 3:
+                        msg += f'{j}. {span}\n'
+                        span_list[f'{i}'].append(f'{j}. {span.strip()}')
+                        doc_map[str(j)] = str(i)
+                        j+=1
+                        docs = docs.replace(span.strip(), f'<highlight_start>{span.strip()}<highlight_end>')
+                i+=1
+        self.pipeline.head['doc_map'] = doc_map
+        self.pipeline.head['docs'] = docs
+        self.pipeline.head['span'] = msg
+        self.pipeline.head['span_list'] = span_list
+        return {'span_list': Prompt.UNABLE,'doc_map': Prompt.UNABLE}
+    
+class Attribute_post_cluster(LLM):
+    noisy = False
+    model_name = 'function'
+    def generate(self, head_prompt: dict = {}, dynamic_prompt: dict = {}):
+        print('f1', head_prompt, dynamic_prompt)
+        span_ls = head_prompt['span_list']
+        doc_map = head_prompt['doc_map']
+        span_list = [item for sublist in head_prompt['span_list'].values() for item in sublist]
+        clusters = eval(one_paragraph(dynamic_prompt['cls'].strip()))
+        self.pipeline.head['clusters'] = clusters
+        def _form(cls):
+            text = ''
+            doc_list = cls['cluster']
+            for doc_num in span_ls.keys():
+                pieces = [str(i) for i in doc_list if doc_map.get(str(i),'None') == doc_num]
+                if pieces:
+                    text += f'Document [{doc_num}]: \n' + '\n'.join([span_list[int(num)-1] for num in pieces])  + '\n'
+
+            return(text)
+        #print([{'span': _form(cls)} for cls in clusters])
+        return [{'span': _form(cls),'span_list': Prompt.UNABLE,'doc_map': Prompt.UNABLE,'clusters':Prompt.UNABLE,'docs':Prompt.UNABLE} for cls in clusters if _form(cls)]
+
+class Ranker(EvalModule):
+    
+    def __init__(self, prompt_maker=None, pipeline=None, self_prompt={}, criteria=None,iterative = True, max_turn = 3, parallel = False, post_processing = lambda x : x, fixed_turn = None) -> None:
+        super().__init__(prompt_maker, pipeline, self_prompt, criteria, iterative = iterative, max_turn = max_turn, parallel = parallel)
+        self.compare = []
+        self.post_processing = post_processing
+        if fixed_turn:
+            self.fixed_turn = fixed_turn
+        else:
+            self.fixed_turn = max_turn
+    def generate(self, head_prompt: dict = {}, dynamic_prompt: dict = {}):
+        
+        self.turns += 1
+        result = {}
+        p_data = {**head_prompt, **self.self_prompt,**dynamic_prompt}
+        for criteria, get_data in self.get_data.items():
+            data_dict = {}
+            for k, v in get_data.items():
+                if v == 'doc_cache':
+                    data_dict[k] = self.pipeline.doc_cache.show_docs()
+                else:
+                    if v in p_data.keys():
+                        data_dict[k] = p_data[v]
+                    else:
+                        data_dict[k] = self.pipeline.current_data[v]
+            eval_func = self.eval_criteria[criteria]
+            data = [data_dict]
+            result[criteria] = eval_func(data)
+        
+        result = sum([value for key, value in result.items()])
+        self.compare.append((result,dynamic_prompt))
+        output = max(self.compare,key = lambda x:x[0])[1]
+        destination = self.send()
+        self.last_message = output
+        if len(self.compare) == self.fixed_turn:
+            self.compare = []
+        if destination in self.conditions:
+            return self.conditions[destination]['post_processing'](output)
+        else: 
+            return self.post_processing(output)
+
+        return {}
+
+    def end_multi(self):
+        self.compare = []
+        return super().end_multi()
+    
+
+
+class AttributingModule(AugmentCluster):
+    model_type = 'attributer'
+    demo ={
+   "selection_instruction":"In this task, you are presented with several documents, which need to be summarized. As an intermediate step, you need to identify salient content within the documents. For each document, copy verbatim the salient spans, and use <SPAN_DELIM> as a delimiter between each consecutive span. IMPORTANT: The output must be of the format Document [<DOC_ID>]: <SPAN_DELIM>-delimited consecutive salient spans. IMPORTANT: Each salient content must be a single consecutive verbatim span from the corresponding passages. IMPORTANT: make sure the total number of copied words (from all documents) is around 200 words, and not more than 900.",
+    "selection_shot":"Document [1]: Cherrapunji Cherrapunji ( with the native name Sohra being more commonly used, and can also be spelled Cherrapunjee or Cherrapunji) is a subdivisional town in the East Khasi Hills district in the Indian state of Meghalaya. It is the traditional capital of aNongkhlaw \"hima\" (Khasi tribal chieftainship constituting a petty state), both known as Sohra or Churra. Cherrapunji has often been credited as being the wettest place on Earth, but for now nearby Mawsynram currently holds that distinction. Cherrapunji still holds the all-time record for the most rainfall in a calendar month for July 1861 and most rain in a year from August 1860 to July 1861, however: it received in\" \nDocument [2]: Radio relay station known as Akashvani Cherrapunji. It broadcasts on FM frequencies. Cherrapunji Cherrapunji (; with the native name Sohra being more commonly used, and can also be spelled Cherrapunjee or Cherrapunji) is a subdivisional town in the East Khasi Hills district in the Indian state of Meghalaya. It is the traditional capital of aNongkhlaw \"hima\" (Khasi tribal chieftainship constituting a petty state), both known as Sohra or Churra. Cherrapunji has often been credited as being the wettest place on Earth, but for now nearby Mawsynram currently holds that distinction. Cherrapunji still holds the all-time record for the most rainfall\" \nDocument [3]: \"Mawsynram Mawsynram () is a village in the East Khasi Hills district of Meghalaya state in north-eastern India, 65 kilometres from Shillong. Mawsynram receives one of the highest rainfalls in India. It is reportedly the wettest place on Earth, with an average annual rainfall of 11,872 mm, but that claim is disputed by Lloró, Colombia, which reported an average yearly rainfall of 12,717 mm between 1952 and 1989 and López de Micay, also in Colombia, which reported an annual 12,892 mm per year between 1960 and 2012. According to the \"Guinness Book of World Records\", Mawsynram received of rainfall in 1985. Mawsynram is located at 25° 18′\" \n\nAnswer: \nDocument [1]: <SPAN_DELIM>Cherrapunji has often been credited as being the wettest place on Earth<SPAN_DELIM> still holds the all-time record for the most rainfall in a calendar month for July 1861 and most rain in a year from August 1860 to July 1861<SPAN_DELIM> \nDocument [2]: <SPAN_DELIM>Cherrapunji has often been credited as being the wettest place on Earth<SPAN_DELIM>still holds the all-time record for the most rainfall<SPAN_DELIM> \nDocument [3]: <SPAN_DELIM>Mawsynram receives one of the highest rainfalls in India <SPAN_DELIM> but that claim is disputed by Lloró, Colombia, which reported an average yearly rainfall of 12,717 mm between 1952 and 1989 <SPAN_DELIM> López de Micay, also in Colombia, which reported an annual 12,892 mm per year between 1960 and 2012. <SPAN_DELIM>",
+    "clustering_instruction":"In this task, you are presented with several passages, where some parts are \"highlighted\" (namely, there are <highlight_start> and <highlight_end> tokens before and after each such span). The goal is to fuse all those highlights into a single summary. As an intermediate step, you need to cluster highlights that can be merged into a sentence (namely, each cluster will be later merged into one sentence). Make sure the clusters are in the same order you would then write the corresponding summary sentences. IMORTANT: make sure there are at most 3 clusters, and no more than 3 highlights per cluster. IMPORTANT: The output must be of the format [\"cluster\":[comma-delimited highlights indices]]",
+    "clustering_shot":"Document [1]: Cherrapunji Cherrapunji ( with the native name Sohra being more commonly used, and can also be spelled Cherrapunjee or Cherrapunji) is a subdivisional town in the East Khasi Hills district in the Indian state of Meghalaya. It is the traditional capital of aNongkhlaw \"hima\" (Khasi tribal chieftainship constituting a petty state), both known as Sohra or Churra. <highlight_start>Cherrapunji has often been credited as being the wettest place on Earth<highlight_end>, but for now nearby Mawsynram currently holds that distinction. Cherrapunji <highlight_start>still holds the all-time record for the most rainfall in a calendar month for July 1861 and most rain in a year from August 1860 to July 1861<highlight_end>, however: it received in\" \nDocument [2]: Radio relay station known as Akashvani Cherrapunji. It broadcasts on FM frequencies. Cherrapunji Cherrapunji (; with the native name Sohra being more commonly used, and can also be spelled Cherrapunjee or Cherrapunji) is a subdivisional town in the East Khasi Hills district in the Indian state of Meghalaya. It is the traditional capital of aNongkhlaw \"hima\" (Khasi tribal chieftainship constituting a petty state), both known as Sohra or Churra. <highlight_start>Cherrapunji has often been credited as being the wettest place on Earth<highlight_end>, but for now nearby Mawsynram currently holds that distinction. <highlight_start>Cherrapunji still holds the all-time record for the most rainfall<highlight_end>\" \nDocument [3]: \"Mawsynram Mawsynram () is a village in the East Khasi Hills district of Meghalaya state in north-eastern India, 65 kilometres from Shillong. <highlight_start>Mawsynram receives one of the highest rainfalls in India<highlight_end>. It is reportedly the wettest place on Earth, with an average annual rainfall of 11,872 mm, <highlight_start>but that claim is disputed by Lloró, Colombia, which reported an average yearly rainfall of 12,717 mm between 1952 and 1989<highlight_end> and <highlight_start>López de Micay, also in Colombia, which reported an annual 12,892 mm per year between 1960 and 2012.<highlight_end> According to the \"Guinness Book of World Records\", Mawsynram received of rainfall in 1985. Mawsynram is located at 25° 18′\" \n\nThe highlighted spans are: \nDocument [1]: 1. Cherrapunji has often been credited as being the wettest place on Earth \n2. still holds the all-time record for the most rainfall in a calendar month for July 1861 and most rain in a year from August 1860 to July 1861 \nDocument [2]: \n3. Cherrapunji has often been credited as being the wettest place on Earth \n4. still holds the all-time record for the most rainfall \nDocument [3]: \n5.  Mawsynram receives one of the highest rainfalls in India  \n6. but that claim is disputed by Lloró, Colombia, which reported an average yearly rainfall of 12,717 mm between 1952 and 1989  \n7. López de Micay, also in Colombia, which reported an annual 12,892 mm per year between 1960 and 2012. \n\nAnswer: \nThe highlighted spans are clustered as follows: \n[{\"cluster\":[6,7]}, {\"cluster\":[5]},{\"cluster\":[1,2]}]",
+    "gen_instruction":"In this task, you are presented with several passages, where some parts are \"highlighted\" (namely, there are <highlight_start> and <highlight_end> tokens before and after each such span). You may also be presented with a prefix of the answer. You job is to generate the next sentence of the answer, that covers all and only the \"highlighted\" spans. Make sure it connects well with the prefix(if eixists), and that it covers all and only the \"highlighted\" spans. Always cite for any factual claim. When citing several search results, use [1][2][3]. Cite at least one document and at most three documents in each sentence. If multiple documents support the sentence, only cite a minimum sufficient subset of the documents.",
+    "gen_shot":"Document [1]: Cherrapunji Cherrapunji ( with the native name Sohra being more commonly used, and can also be spelled Cherrapunjee or Cherrapunji) is a subdivisional town in the East Khasi Hills district in the Indian state of Meghalaya. It is the traditional capital of aNongkhlaw \"hima\" (Khasi tribal chieftainship constituting a petty state), both known as Sohra or Churra. <highlight_start>Cherrapunji has often been credited as being the wettest place on Earth<highlight_end>, but for now nearby Mawsynram currently holds that distinction. Cherrapunji <highlight_start>still holds the all-time record for the most rainfall in a calendar month for July 1861 and most rain in a year from August 1860 to July 1861<highlight_end>, however: it received in\" \nDocument [2]: Radio relay station known as Akashvani Cherrapunji. It broadcasts on FM frequencies. Cherrapunji Cherrapunji (; with the native name Sohra being more commonly used, and can also be spelled Cherrapunjee or Cherrapunji) is a subdivisional town in the East Khasi Hills district in the Indian state of Meghalaya. It is the traditional capital of aNongkhlaw \"hima\" (Khasi tribal chieftainship constituting a petty state), both known as Sohra or Churra. <highlight_start>Cherrapunji has often been credited as being the wettest place on Earth<highlight_end>, but for now nearby Mawsynram currently holds that distinction. <highlight_start>Cherrapunji still holds the all-time record for the most rainfall<highlight_end>\" \nDocument [3]: \"Mawsynram Mawsynram () is a village in the East Khasi Hills district of Meghalaya state in north-eastern India, 65 kilometres from Shillong. <highlight_start>Mawsynram receives one of the highest rainfalls in India<highlight_end>. It is reportedly the wettest place on Earth, with an average annual rainfall of 11,872 mm, <highlight_start>but that claim is disputed by Lloró, Colombia, which reported an average yearly rainfall of 12,717 mm between 1952 and 1989<highlight_end> and <highlight_start>López de Micay, also in Colombia, which reported an annual 12,892 mm per year between 1960 and 2012.<highlight_end> According to the \"Guinness Book of World Records\", Mawsynram received of rainfall in 1985. Mawsynram is located at 25° 18′\" \n\nPrefix: Several places on Earth claim to be the most rainy, such as Lloró, Colombia, which reported an average annual rainfall of 12,717 mm between 1952 and 1989, and López de Micay, Colombia, which reported an annual 12,892 mm between 1960 and 2012 [3].  \n\nThe highlighted spans are: \nDocument [3]: \n5.  Mawsynram receives one of the highest rainfalls in India \n\nAnswer: \nThe next sentence is: \nHowever, the official record is held by Mawsynram, India with an average annual rainfall of 11,872 mm [3]."
+    }
+    PARA_SEP = '\n\n'
+    selection_shot = demo['selection_instruction'] + PARA_SEP + demo['selection_shot'] + PARA_SEP
+    cls_shot = demo['clustering_instruction'] + PARA_SEP + demo['clustering_shot'] + PARA_SEP
+    gen_shot = demo['gen_instruction'] + PARA_SEP + demo['gen_shot'] + PARA_SEP
+    prompt = Prompt(template='<shot><INST><question><docs><prefix><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',
+                            'add':'Answer: \n{add}'
+                            })
+    def __init__(self, model) -> None:
+        module_list = []
+        select = LLM(model = model, prompt_maker = self.prompt, self_prompt={'INST':self.demo['selection_instruction'],'shot':self.selection_shot,'add':''}, post_processing=make_as('span'),noisy= False)
+        post_select = Attribute_post_select()
+        clustering =  LLM(model = model, prompt_maker = self.prompt, self_prompt={'INST':self.demo['clustering_instruction'],'shot':self.cls_shot, 'add':'The highlighted spans are clustered as follows:'},share_model_with=select, post_processing=make_as('cls'),noisy=False)
+        post_cls = Attribute_post_cluster()
+        module_list = [select,post_select,clustering,post_cls]
+        super().__init__(module_list)
+
+MODEL_TYPE_MAPPING = {
+    'retrieve': Retriever,
+    'eval': EvalModule,
+    'simplify': CitationSimplyfier, 
+    'verify': Verifier,
+    'rank': Ranker,
+    'attributing': AttributingModule
+}

citekit/evaluator/evaluator.py

@@ -0,0 +1,1118 @@
+from nltk import sent_tokenize
+import nltk
+nltk.download('punkt')
+import re 
+import random
+import transformers
+import numpy as np
+from citekit.utils.utils import *
+from rouge import Rouge
+import torch
+from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
+import copy
+import torch
+from tqdm import tqdm
+import sys
+import logging
+import random
+from itertools import product,combinations
+import time
+import logging
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.INFO)
+
+PIPELINE_OUTPUT = 'output'
+PIPELINE_DOC_CACHE = 'doc_cache'
+
+global autoais_model, autoais_tokenizer
+autoais_model = None
+autoais_tokenizer = None
+get_docs_by_index = lambda i,docs: docs[i] if i < len(docs) else None 
+ais_LLM = None
+
+QA_MODEL = "gaotianyu1350/roberta-large-squad"
+AUTOAIS_MODEL = "google/t5_xxl_true_nli_mixture"
+AUTOAIS_MODEL_ABSOLUTE = "/mnt/usercache/huggingface/t5_xxl_true_nli_mixture"
+
+def get_cite(sent):
+    return re.sub(r"\[\d+", "", re.sub(r" \[\d+", "", sent)).replace(" |", "").replace("]", ""),[int(r[1:]) - 1 for r in re.findall(r"\[\d+", sent)]
+
+
+def entail(premise, claim):
+
+    """
+    Run inference for assessing AIS between a premise and hypothesis.
+    Adapted from https://github.com/google-research-datasets/Attributed-QA/blob/main/evaluation.py
+    """
+    global autoais_model, autoais_tokenizer
+    input_text = "premise: {} hypothesis: {}".format(premise, claim)
+    input_ids = autoais_tokenizer(input_text, return_tensors="pt").input_ids.to(autoais_model.device)
+    with torch.inference_mode():
+        outputs = autoais_model.generate(input_ids, max_new_tokens=10)
+    result = autoais_tokenizer.decode(outputs[0], skip_special_tokens=True)
+    inference = 1 if result == "1" else 0
+    return inference
+
+def load_auto_ais():
+    global autoais_model, autoais_tokenizer
+    print('Initializing eval model for citation precision and recall...') 
+    try:
+        autoais_model = AutoModelForSeq2SeqLM.from_pretrained(AUTOAIS_MODEL, torch_dtype=torch.bfloat16, device_map="auto")
+        autoais_tokenizer = AutoTokenizer.from_pretrained(AUTOAIS_MODEL, use_fast=False)
+        
+    except:
+        print('Unable to load model from hub, trying to load from local path...')
+        autoais_model = AutoModelForSeq2SeqLM.from_pretrained(AUTOAIS_MODEL_ABSOLUTE, torch_dtype=torch.bfloat16, device_map="auto")
+        autoais_tokenizer = AutoTokenizer.from_pretrained(AUTOAIS_MODEL_ABSOLUTE, use_fast=False)
+    print('Done!')
+
+def compute_mauve(data):
+    """Compute Mauve score."""
+
+    logger.info("Computing MAUVE...")
+    human_data = []
+    model_data = []
+    for item in data:
+        # Remove ending punctuations
+        # Remove any new lines
+        # Truncate by 100 words
+        human_data.append(
+            ' '.join((item['question'] + " " + item['answer'].strip()).split()[:100]).rstrip(string.punctuation))
+        model_data.append(
+            ' '.join((item['question'] + " " + item['output'].strip()).split()[:100]).rstrip(string.punctuation))
+
+    import mauve
+    out = mauve.compute_mauve(
+        p_text=human_data,
+        q_text=model_data,
+        device_id=0,
+        max_text_length=512,
+        verbose=True,
+        batch_size=8,
+        featurize_model_name="gpt2-large"
+    )
+    return out.mauve * 100
+
+
+def compute_rouge_l(data):
+    total = len(data)
+    res = {
+                "r": 0.0,
+                "p": 0.0,
+                "f": 0.0
+            }
+    for item in data:
+        if item['output'] and item['answer']:
+            rouge = Rouge()
+            scores = rouge.get_scores(item['output'], item['answer'])
+            res['r'] += scores[0]['rouge-l']['r']
+            res['p'] += scores[0]['rouge-l']['p']
+            res['f'] += scores[0]['rouge-l']['f']
+        else:
+            print('Warning: no hypothesis or references')
+    res['r'] /= total
+    res['p'] /= total
+    res['f'] /= total
+
+    return res
+    
+def compute_qa(question, output, short_answers, qa_pipeline=None):
+    """Compute QA-based accuracy.
+    Args:
+        
+    Returns:
+        QA metrics (QA-EM, QA-F1, QA-Hit)
+    """
+
+    # Load model
+    if not qa_pipeline:
+        qa_pipeline = transformers.pipeline("question-answering", model=QA_MODEL, device='mps')
+
+    # Get prediction
+    em, f1, bins = 0,0,0
+    context = output if len(output) > 0 else " "
+    result = qa_pipeline(question=question, context=context, handle_impossible_answer=True)
+    loc_counter, loc_em, loc_f1 = 0, 0, 0
+    print(result)
+    prediction = result["answer"]
+
+    loc_em = max([compute_exact(a, prediction) for a in short_answers])
+    loc_f1 = max([compute_f1(a, prediction) for a in short_answers])
+    loc_counter += 1
+
+    em= loc_em / loc_counter
+    f1= loc_f1 / loc_counter 
+    bins = int(loc_em == loc_counter)
+    return em, f1, bins
+
+def compute_qa(data):
+    """Compute QA-based accuracy.
+    Args:
+        data: requires filed `qa_pairs/short_answers` and `output`
+    Returns:
+        QA metrics (QA-EM, QA-F1, QA-Hit)
+    """
+
+    if 'qa_pairs' not in data[0] or data[0]['qa_pairs'] is None:
+        #logger.warn("Warning: no QA pairs found in data")
+        return {
+            'QA-EM': 0,
+            'QA-F1': 0,
+            'QA-Hit': 0,
+        }
+
+    # Load model
+    #logger.info("Loading the RoBERTa-large SQuAD model for QA-based accuracy...")
+    global qa_pipeline
+    if not qa_pipeline:
+        qa_pipeline = transformers.pipeline("question-answering", model=QA_MODEL)
+    #logger.info("Done")
+
+    # Get prediction
+    #logger.info("Computing the QA-based accuracy...")
+    em, f1, bins = [], [], []
+    for item in tqdm(data):
+        question = [qa_pair['question'] for qa_pair in item['qa_pairs']]
+        context = item['output'] if len(item['output']) > 0 else " "
+        results = qa_pipeline(question=question, context=context, handle_impossible_answer=True)
+        loc_counter, loc_em, loc_f1 = 0, 0, 0
+
+        for idx, res in enumerate(results):
+            answers = item["qa_pairs"][idx]["short_answers"]
+            prediction = res["answer"]
+
+            loc_em += max([compute_exact(a, prediction) for a in answers])
+            loc_f1 += max([compute_f1(a, prediction) for a in answers])
+            loc_counter += 1
+
+        em.append(loc_em / loc_counter)
+        f1.append(loc_f1 / loc_counter)
+        bins.append(loc_em == loc_counter)
+
+    return {
+        'QA-EM': 100 * np.mean(em),
+        'QA-F1': 100 * np.mean(f1),
+        'QA-Hit': 100 * np.mean(bins)
+    }
+
+
+def cite_pr(sent_with_cite, docs = None, get_docs = get_docs_by_index, get_cite = get_cite, max_cite= None,rich_return = False):
+    """
+    : sent_with_cite: ONE sentence with citation like [1][2][3]
+    : get_docs: by default like [1][2], get ids
+    : docs: List, all the COMPLETE documents with TITLE
+
+    : return 
+        number of citations, integer
+        recall (0 or 1)
+        precision (number of relevent documents)
+
+        optional;
+            multi_cite
+            mcite_support
+            mcite_overcite
+    """
+    if rich_return:
+        raise NotImplementedError
+
+    result = {'num_cites': 0,'recall':0,'precision':0,'multi_cite':0,'mcite_support' :0,'mcite_overcite':0}
+    sent, cites= get_cite(sent_with_cite)
+
+    if not cites:
+        return (0, 0, 0) if not rich_return else result # no citations
+    if max_cite:
+        cites = cites[:max_cite]
+    num_cites = len(cites)
+    result['num_cites'] = num_cites
+
+    refs = [get_docs(cite, docs) for cite in cites]
+    if None in refs:
+        return (num_cites, 0, 0) if not rich_return else result# wrong citation(s)
+    
+    # recall
+    recall = entail(premise=''.join(refs),claim=sent)
+    result['recall'] = recall
+
+    # precision
+    precision = 0
+    if num_cites == 1:
+        precision = recall
+    else:
+        for idx, ref in enumerate(refs):
+            if entail(premise=ref,claim=sent):
+                precision += 1
+            else:
+                if not entail(premise=''.join([refs[i] for i in range(len(refs)) if i != idx]), claim = sent):
+                    precision += 1
+                elif recall:
+                    result['mcite_overcite'] = 1
+    result['precision'] = precision
+    
+    #other 
+    if num_cites > 1:
+        result['multi_cite'] = 1
+        if recall:
+            result['mcite_support'] = 1
+    
+
+    return (num_cites, recall, precision) if not rich_return else result
+
+
+def cite_pr_answer(answer, docs = None, get_docs = get_docs_by_index, get_cite = get_cite, max_cite= None,rich_return = False):
+    epsilon = 1e-8
+    num_c = 0
+    recall = 0
+    precision = 0
+    sents = sent_tokenize(answer)
+    for sent in sents:
+        c,r,p = cite_pr(sent,get_docs=get_docs,docs=docs,get_cite=get_cite,max_cite=max_cite,rich_return=rich_return)
+        num_c += c
+        recall += r
+        precision += p
+    # diveded by Zero!
+    return recall/(len(sents)+ epsilon), precision/(num_c+epsilon)
+
+
+def _run_nli_autoais(passage, claim, test = False):
+    """
+    Run inference for assessing AIS between a premise and hypothesis.
+    Adapted from https://github.com/google-research-datasets/Attributed-QA/blob/main/evaluation.py
+    """
+    if not test:
+        global autoais_model, autoais_tokenizer
+        if not autoais_model:
+            load_auto_ais()
+        input_text = "premise: {} hypothesis: {}".format(passage, claim)
+        input_ids = autoais_tokenizer(input_text, return_tensors="pt").input_ids.to(autoais_model.device)
+        with torch.inference_mode():
+            outputs = autoais_model.generate(input_ids, max_new_tokens=10)
+        result = autoais_tokenizer.decode(outputs[0], skip_special_tokens=True)
+        inference = 1 if result == "1" else 0
+        return inference
+    else:
+        res = random.randint(0,1)
+
+    return res
+
+def _run_llm_autoais(passage, claim):
+    global ais_LLM
+    assert(ais_LLM)
+    return int(ais_LLM.generate(premise = passage, claim = claim))
+
+def test_compute_autoais(data):
+    print(data[0]['docs'][:5])
+    print(data[0]['output'][:5])
+    return {
+        "citation_rec": random.randint(0,100),
+        "citation_prec": random.randint(0,100),
+    }
+
+def compute_autoais(data,
+                    decontext=False,
+                    concat=False,
+                    qampari=False,
+                    at_most_sents = 3,
+                    at_most_citations=3,
+                    entail_function = _run_nli_autoais):
+    """
+    Compute AutoAIS score.
+
+    Args:
+        data: requires field `output` and `docs`
+              - docs should be a list of items with fields `title` and `text` (or `phrase` and `sent` for QA-extracted docs)
+        citation: check citations and use the corresponding references.
+        decontext: decontextualize the output
+    """
+
+    global autoais_model, autoais_tokenizer
+
+
+    ais_scores = []
+    ais_scores_prec = []
+
+    sent_total = 0
+    sent_mcite = 0
+    sent_mcite_support = 0
+    sent_mcite_overcite = 0
+    autoais_log = []
+    for item in tqdm(data):
+        # Get sentences by using NLTK
+        if qampari:
+            print('now qampari...')
+            sents = [item['question'] + " " + x.strip() for x in
+                     item['output'].rstrip().rstrip(".").rstrip(",").split(",")]
+        else:
+            sents = sent_tokenize(item['output'])[:at_most_sents]
+        if len(sents) == 0:
+            ais_scores.append(0.0)
+            ais_scores_prec.append(0.0)  # len(sents))
+            continue
+
+        target_sents = [remove_citations(sent).strip() for sent in sents]
+
+        entail = 0
+        entail_prec = 0
+        total_citations = 0
+        for sent_id, sent in enumerate(sents):
+            target_sent = target_sents[sent_id]  # Citation removed and (if opted for) decontextualized
+            joint_entail = -1  # Undecided
+
+            # Find references
+            #ref = [int(r[1:]) - 1 for r in re.findall(r"\[\d+", sent)]  # In text citation id starts from 1
+            matches = re.findall(r"\[(\d+(?:,\s*\d+)*)\]", sent)
+            ref = [int(num)-1 for match in matches for num in match.replace(' ', '').split(',')]
+            if len(ref) == 0:
+                # No citations
+                joint_entail = 0
+            elif any([ref_id >= len(item['docs']) for ref_id in ref]):
+                # Citations out of range
+                joint_entail = 0
+            else:
+                if at_most_citations is not None:
+                    ref = ref[:at_most_citations]
+                total_citations += len(ref)
+                joint_passage = '\n'.join([(item['docs'][psgs_id]) for psgs_id in ref])
+
+            # If not directly rejected by citation format error, calculate the recall score
+            if joint_entail == -1:
+                joint_entail = entail_function(joint_passage, target_sent)
+                autoais_log.append({
+                    #"question": item['question'],
+                    "output": item['output'],
+                    "claim": sent,
+                    "passage": [joint_passage],
+                    "model_type": "NLI",
+                    "model_output": joint_entail,
+                })
+
+            entail += joint_entail
+            if len(ref) > 1:
+                sent_mcite += 1
+
+            # calculate the precision score if applicable
+            if joint_entail and len(ref) > 1:
+                sent_mcite_support += 1
+                # Precision check: did the model cite any unnecessary documents?
+                for psgs_id in ref:
+                    # condition A
+                    passage = item['docs'][psgs_id]
+                    nli_result = entail_function(passage, target_sent)
+
+                    # condition B
+                    if not nli_result:
+                        subset_exclude = copy.deepcopy(ref)
+                        subset_exclude.remove(psgs_id)
+                        passage = '\n'.join([item['docs'][pid] for pid in subset_exclude])
+                        nli_result =entail_function(passage, target_sent)
+                        if nli_result:  # psgs_id is not necessary
+                            flag = 0
+                            sent_mcite_overcite += 1
+                        else:
+                            entail_prec += 1
+                    else:
+                        entail_prec += 1
+            else:
+                entail_prec += joint_entail
+        sent_total += len(sents)
+        ais_scores.append(entail / len(sents))
+        ais_scores_prec.append(entail_prec / total_citations if total_citations > 0 else 0)  # len(sents))
+
+    if sent_mcite > 0 and sent_mcite_support > 0:
+        print(
+            "Among all sentences, %.2f%% have multiple citations, among which %.2f%% are supported by the joint set, among which %.2f%% overcite." % (
+                100 * sent_mcite / sent_total,
+                100 * sent_mcite_support / sent_mcite,
+                100 * sent_mcite_overcite / sent_mcite_support
+            ))
+
+    return {
+        "citation_rec": 100 * np.mean(ais_scores),
+        "citation_prec": 100 * np.mean(ais_scores_prec),
+    }
+
+def compute_claims_test(data):
+    print(data[0]['claims'])
+    print(data[0][PIPELINE_OUTPUT])
+    return random.randint(1,100)
+
+def compute_claims(data):
+    global autoais_model, autoais_tokenizer
+    if autoais_model is None:
+        #logger.info("Loading AutoAIS model...")
+        # autoais_model = AutoModelForSeq2SeqLM.from_pretrained(AUTOAIS_MODEL, torch_dtype=torch.bfloat16, max_memory=get_max_memory(), device_map="auto")
+        autoais_model = AutoModelForSeq2SeqLM.from_pretrained(AUTOAIS_MODEL, torch_dtype=torch.bfloat16,
+                                                              device_map="auto")
+        # autoais_model = AutoModelForSeq2SeqLM.from_pretrained(AUTOAIS_MODEL, torch_dtype=torch.bfloat16, max_memory=get_max_memory(), device_map="auto",offload_folder= "/data/hongbang/zsf/projects/ALCE/ALCE/model/t5_xxl_true_nli_mixture/offload1")
+        autoais_tokenizer = AutoTokenizer.from_pretrained(AUTOAIS_MODEL, use_fast=False)
+    #logger.info("Computing claims...")
+    scores = []
+    for item in tqdm(data):
+        normalized_output = remove_citations(item['output'])
+        entail = 0
+        claims = item["claims"]
+        for claim in claims:
+            entail += _run_nli_autoais(normalized_output, claim)
+        scores.append(entail / len(claims))
+    return 100 * np.mean(scores)
+
+
+#citation appropriateness
+def check_if_citations_needed(passages, answer, grain):
+
+    def _format_document(doc):
+        """Format document for AutoAIS.
+
+        if "sent" in doc:
+            # QA-extracted docs
+            return "Title: %s\n%s" % (doc['title'], doc['sent'])
+        else:
+            return "Title: %s\n%s" % (doc['title'], doc['text'])
+        """
+        return doc
+
+    global autoais_model, autoais_tokenizer
+    if autoais_model is None and False:
+        #logger.info("Loading AutoAIS model...")
+        # autoais_model = AutoModelForSeq2SeqLM.from_pretrained(AUTOAIS_MODEL, torch_dtype=torch.bfloat16, max_memory=get_max_memory(), device_map="auto")
+        autoais_model = AutoModelForSeq2SeqLM.from_pretrained(AUTOAIS_MODEL, torch_dtype=torch.bfloat16,
+                                                              device_map="auto")
+        # autoais_model = AutoModelForSeq2SeqLM.from_pretrained(AUTOAIS_MODEL, torch_dtype=torch.bfloat16, max_memory=get_max_memory(), device_map="auto",offload_folder= "/data/hongbang/zsf/projects/ALCE/ALCE/model/t5_xxl_true_nli_mixture/offload1")
+        autoais_tokenizer = AutoTokenizer.from_pretrained(AUTOAIS_MODEL, use_fast=False)
+
+    if grain=="over_fine" or grain=="more_over_fine":
+        num_passages = len(passages)
+        passages_per_chunk = num_passages // 5  # Divide passages evenly into 5 chunks
+        remainder = num_passages % 5  # Handle remaining passages
+        passages_five=[]
+        start_idx = 0
+        for i in range(5):
+            end_idx = start_idx + passages_per_chunk
+            if remainder > 0:
+                end_idx += 1
+                remainder -= 1
+            chunk_passages = passages[start_idx:end_idx]
+            passages_five.append('\n'.join([_format_document(p) for p in chunk_passages]))
+            start_idx = end_idx
+        passages=passages_five
+        combinations_3 = combinations(passages, 3)  # 获取所有三个passage的组合
+        for combination in combinations_3:
+            joint_passage = '\n'.join(
+                [passage for passage in combination])  # 将三个passage连接为一个字符串，并保留格式
+            entail = _run_nli_autoais(joint_passage, answer)
+            if entail == 1:
+                return 1
+        return 0
+
+    else:
+        if len(passages)>=3:#正常粒度
+            combinations_3 = combinations(passages, 3)
+            for combination in combinations_3:
+                joint_passage = '\n'.join(
+                    [_format_document(passage) for passage in combination])
+                entail = _run_nli_autoais(joint_passage, answer)
+                if entail == 1:
+                    return 1
+            return 0
+        else:#粗粒度
+            joint_passage = '\n'.join(
+                [_format_document(passage) for passage in passages])
+            entail = _run_nli_autoais(joint_passage, answer)
+            if entail == 1:
+                return 1
+            else:
+                return 0
+
+
+#citaion granularity
+def find_permutations(n, m):
+    '''
+    :param n:  最大数量总和
+    :param m: 位长度
+    :return:
+    '''
+    # Generate all possible sequences of length m
+    all_sequences = list(product(range(n + 1), repeat=m))
+    #print('all_sequences', all_sequences)
+
+    # Filter sequences where the sum of digits equals n
+    valid_sequences = [seq for seq in all_sequences if sum(seq) == n]
+    return valid_sequences
+
+
+def get_subspans(list_span, span_count):
+    list_subspan = []
+    for i in range(0, len(list_span) - span_count + 1):
+        list_subspan.append(list_span[i: i + span_count])
+
+    return list_subspan
+
+
+def get_all_span_comb(list_list_span, target_span_count=-1):
+    if target_span_count == -1: # 所有子集
+        max_span_count = len(sum(list_list_span, []))
+        doc_count = len(list_list_span)
+        list_span_comb_all = []
+        for span_count in range(1, max_span_count + 1):
+            list_comb = find_permutations(span_count, doc_count)#给定数量的子串在文本中的所有可能组合
+
+            list_span_comb = [] # 最终当前长度的所有可能组合
+            for comb in list_comb:
+                list_list_subspan = []
+
+                for idx_doc, span_count_doc in enumerate(comb):
+                    list_subspan = get_subspans(list_list_span[idx_doc], span_count_doc)
+                    if len(list_subspan) == 0:
+                        list_list_subspan = None
+                        break
+                    list_list_subspan.append(list_subspan)
+
+                if list_list_subspan:
+                    list_span_comb_cur = [sum(list(combination), [])  for combination in product(*list_list_subspan)]
+                    list_span_comb_cur = list(set([tuple(span_comb) for span_comb in list_span_comb_cur]))
+
+                    list_span_comb += list_span_comb_cur
+            list_span_comb_all += list_span_comb
+        list_span_comb_all = set(list_span_comb_all)
+    else: # 当前长度的组合
+        doc_count = len(list_list_span)
+        list_comb = find_permutations(target_span_count, doc_count)
+
+        list_span_comb = []  # 最终当前长度的所有可能组合
+        for comb in list_comb:
+            list_list_subspan = []
+
+            for idx_doc, span_count_doc in enumerate(comb):
+                list_subspan = get_subspans(list_list_span[idx_doc], span_count_doc)
+                if len(list_subspan) == 0:
+                    list_list_subspan = None
+                    break
+                list_list_subspan.append(list_subspan)
+
+            if list_list_subspan:
+                list_span_comb_cur = [combination for combination in product(*list_list_subspan)]
+                for idx in range(len(list_span_comb_cur)):
+                    list_span_comb_cur[idx] = tuple([tuple(span_comb) for span_comb in list_span_comb_cur[idx]])
+
+                list_span_comb += list_span_comb_cur
+        list_span_comb_all = list_span_comb
+        list_span_comb_all = set(list_span_comb_all)
+    return list_span_comb_all
+
+
+def run_converge_2(list_list_span=None, sentence=None):
+    '''
+    基于假设：更长的text不能蕴含，则其任何子串都不能蕴含
+    span数量递减（提供更多的剪枝选项）
+    最终gold可能有一个span的误差
+    '''
+    ######
+    #print('origin nli count', len(get_all_span_comb(list_list_span, target_span_count=-1)))#给定文本的所有可能的子串组合
+    max_span_count = len(sum(list_list_span, [])) # span总数
+
+    set_comb_hash = set([])
+
+    ### span数量二分
+    nli_count = 0
+    skip_count = 0
+    list_list_span_gold = copy.copy(list_list_span) # 当前能够精准蕴含的span
+
+    span_count_min, span_count_max = 1, max_span_count
+    start_time=time.time()
+    timeout=300
+    while span_count_min < span_count_max:#每次迭代中不断寻找更小的子串组合
+        span_count_cur = span_count_max - 1
+        flag_find = False
+        if time.time() - start_time > timeout:
+            print('timeout!')
+            list_list_span_gold=[]
+            break
+        ### 存在可蕴含，继续找更少的span
+        ### 不存在可蕴含，继续找更多的span
+        # 长度为span_count_max - 1的所有可能的子串组合
+        set_comb_cur = get_all_span_comb(list_list_span, target_span_count=span_count_cur)
+
+        list_comb_cur = list(set_comb_cur)
+        random.shuffle(list_comb_cur)
+        for comb in list_comb_cur:
+            list_list_span_cur = [list(t) for t in comb]
+            list_span_cur = sum(list_list_span_cur, [])
+            str_text = ' '.join(list_span_cur) # TODO: 统一字符串化的方式
+
+            if hash(str_text) in set_comb_hash:
+                skip_count += 1
+                continue
+
+            #### ⚠️ 注意在这里替换nli函数
+            nli_label = _run_nli_autoais(str_text, sentence) # TODO: nli label function
+            nli_count += 1
+
+            if nli_label == 1: # 只要存在可蕴含，直接继续找更少的span
+                list_list_span_gold = copy.copy(list_list_span_cur)
+                span_count_max = span_count_cur#更新span数量上限
+                flag_find = True
+                # print(f"find nli!, nli_count: {nli_count}, skip_count: {skip_count}, len(set_comb_hash): {len(set_comb_hash)}", )
+                break
+            else: # 不能蕴含，剪枝所有子集
+                set_comb_cur_del = get_all_span_comb(list_list_span_cur, target_span_count=-1)
+                set_comb_hash_cur = set([hash(' '.join(list(tuple_comb_))) for tuple_comb_ in set_comb_cur_del]) # TODO: 统一字符串化的方式
+
+                set_comb_hash |= set_comb_hash_cur
+        if flag_find == False:
+            print(f"CAN'T find nli!, nli_count: {nli_count}, skip_count: {skip_count}, len(set_comb_hash): {len(set_comb_hash)}", )
+            break
+    span_count_gold = span_count_max # gold的span数量
+    print('len(set_comb_del)', len(set_comb_hash))
+    print('nli_count', nli_count, 'skip_count', skip_count, 'span_count_gold', span_count_gold)
+    return list_list_span_gold
+
+
+def compute_autoais_grained(data,
+                    at_most_citations=3,method='ALCE',grain='default'):
+
+    """
+    Compute AutoAIS score.
+
+    Args:
+        data: requires field `output` and `docs`
+              - docs should be a list of items with fields `title` and `text` (or `phrase` and `sent` for QA-extracted docs)
+        citation: check citations and use the corresponding references.
+        decontext: decontextualize the output
+    """
+    global autoais_model, autoais_tokenizer
+    if autoais_model is None and False:
+        #logger.info("Loading AutoAIS model...")
+        # autoais_model = AutoModelForSeq2SeqLM.from_pretrained(AUTOAIS_MODEL, torch_dtype=torch.bfloat16, max_memory=get_max_memory(), device_map="auto")
+        autoais_model = AutoModelForSeq2SeqLM.from_pretrained(AUTOAIS_MODEL, torch_dtype=torch.bfloat16,
+                                                              device_map="auto")
+        # autoais_model = AutoModelForSeq2SeqLM.from_pretrained(AUTOAIS_MODEL, torch_dtype=torch.bfloat16, max_memory=get_max_memory(), device_map="auto",offload_folder= "/data/hongbang/zsf/projects/ALCE/ALCE/model/t5_xxl_true_nli_mixture/offload1")
+        autoais_tokenizer = AutoTokenizer.from_pretrained(AUTOAIS_MODEL, use_fast=False)
+    def _format_document(doc):
+
+        """Format document for AutoAIS."""
+        if isinstance(doc, dict):
+            if "sent" in doc:
+                # QA-extracted docs
+                return "Title: %s\n%s" % (doc['title'], doc['sent'])
+            else:
+                return "Title: %s\n%s" % (doc['title'], doc['text'])
+        elif isinstance(doc,str):
+            return doc
+
+
+    #logger.info(f"Running AutoAIS...")
+
+    ais_scores_need = []  # 是否需要引用
+    ais_scores = []  # quote_recall
+    ais_doc_scores=[]#doc_recall
+
+    sent_total = 0
+
+    autoais_log = []
+    granularity_list = []
+    skipped =0
+    for item in tqdm(data):
+        output = item['output']
+
+        if method=='baseline':
+            model_answer=item['output_parse']['answer']
+            answer = ''
+            reference = {}
+            span_contents = {}
+            if not model_answer["text"].endswith("."):
+                model_answer["text"] += "."
+            answer += " " + model_answer["text"]
+            spans = model_answer['reference']
+            for span in spans:
+                match = re.match(r'^(\d+)\.', span)
+                if match:
+                    span_number = match.group(1)
+                    span_content = span.split('. ', 1)[1].strip()  # 获取1. 后面的内容
+                    span_contents[span_number] = span_content
+            reference.update(span_contents)
+
+            item['output_answer'] = answer.strip()
+            item['output_ref_span'] = reference
+            output = item['output_answer']
+
+        elif method=='ALCE':
+            # 匹配 According to Document
+            pattern_doc = r"According to Document \[(\d+)\]"
+            # 匹配 (Title: Godfrey Chitalu)
+            pattern_title = r"\(Title: [^\)]+\)"
+
+            output = re.sub(pattern_doc, r"[\1]", output)
+            output = re.sub(pattern_title, "", output)
+            output=output.strip().split("\n")[0]
+            output=output.replace("<|im_end|>", "")
+        # Get sentences by using NLTK
+        sents = sent_tokenize(output)[:3]
+        if len(sents) == 0:
+            continue
+
+        target_sents = [remove_citations(sent).strip() for sent in sents]
+        output_ref_span = item.get('output_ref_span', {})
+        # sent_joint_passage = '\n'.join([_format_document(doc) for doc in item['docs']])
+
+        entail = 0
+        entail_doc=0
+        total_citations = 0
+        need_citations_sentences = 0  # 一个回答中需要引用的句子数量
+        correct_predictions = 0  # 新增：记录正确的预测是否需要引用的子句数量
+
+        for sent_id, sent in enumerate(sents):
+            target_sent = target_sents[sent_id]  # Citation removed and (if opted for) decontextualized
+            joint_entail = -1  # Undecided
+            joint_doc_entail=-1
+
+            # 1. appropriatness
+            # 每句话是否需要引用
+            need_citations = check_if_citations_needed(item['docs'], target_sent,grain)
+
+
+            if method=='baseline':
+                # Find references number
+                ref_mark = [int(r[1:]) for r in re.findall(r"\{\d+", sent)]
+                # 引用的span（拼��）match document
+                ref, ref_span = match_document(ref_mark, output_ref_span)
+                #logger.info(f"For `{target_sent}`, find citations {ref}")
+                ref_id = [x -1 for x in ref]
+                processed_refs = set()
+                ref_passage = []
+                for psgs_id in ref_id:
+                    if 0 <= psgs_id < len(item['docs']) and psgs_id not in processed_refs:
+                        ref_passage.append(_format_document(item['docs'][psgs_id]))
+                        processed_refs.add(psgs_id)
+                    elif psgs_id in processed_refs:
+                        print("Warning: psgs_id already processed:", psgs_id + 1)
+                    else:
+                        print("Error: psgs_id out of range:", psgs_id+1)
+
+                joint_span = '\n'.join(ref_span)
+                joint_passage = '\n'.join(ref_passage)
+
+            elif method=='ALCE':
+                ref = list(set([int(r[1:]) for r in re.findall(r"\[\d+", sent)]))
+                #logger.info(f"For `{target_sent}`, find citations {ref}")
+                ref_id=list(set([int(r[1:])-1 for r in re.findall(r"\[\d+", sent)]))
+                processed_refs = set()
+                ref_passage = []
+                for psgs_id in ref_id:
+                    if 0 <= psgs_id < len(item['docs']) and psgs_id not in processed_refs:
+                        ref_passage.append(_format_document(item['docs'][psgs_id]))
+                        processed_refs.add(psgs_id)
+                    elif psgs_id in processed_refs:
+                        print("Warning: psgs_id already processed:", psgs_id+1)
+                    else:
+                        print("Error: psgs_id out of range:", psgs_id+1)
+                ref_span=ref_passage
+                joint_passage = '\n'.join(ref_passage)
+                joint_span=joint_passage
+
+
+            autoais_log.append({
+                "question": item['question'],
+                "output_answer": item['output'],
+                "docs": item['docs'],
+                "claim": {
+                    "sentence": sent,
+                    "if_citations_needed": need_citations,
+                    "has_reference": ref,
+                    "doc_recall": None,
+                    "quote_recall": None,
+                    "granularity_score":None,
+                    "granularity_span":None
+                }
+            })
+
+            if len(ref) == 0:
+                # No citations
+                joint_entail = 0
+                joint_doc_entail=0
+            elif any([ref_id > len(item['docs']) for ref_id in ref]):
+                # Citations out of range
+                joint_entail = 0
+                joint_doc_entail=0
+            else:
+                if at_most_citations is not None:
+                    ref = ref[:at_most_citations]
+                total_citations += len(ref)
+
+            # 更新正确预测是否需要引用的数量
+            if_citations_needed = autoais_log[-1]["claim"]["if_citations_needed"]
+            has_reference = autoais_log[-1]["claim"]["has_reference"]
+            if (if_citations_needed == 1 and has_reference) or (if_citations_needed == 0 and not has_reference):
+                correct_predictions += 1
+            #logger.info("citation appropriateness finished")
+
+            # 2. 在需要引用的情况下才计算citation correctness
+            if need_citations and has_reference:#需要引用且引用了才考虑后两个指标
+                start_time = time.time()
+                need_citations_sentences += 1
+                # 2.(1):quote_corr
+                # If not directly rejected by citation format error, calculate the recall score
+                if joint_entail == -1:
+                    # φ(premise, hypothesis)判断所有引用span的拼接是否entail模型的回答output
+                    joint_entail = _run_nli_autoais(joint_span, target_sent)
+                entail += joint_entail
+                autoais_log[-1]["claim"]["quote_recall"] = joint_entail
+                #logger.info(f"citation recall finished, recall is {joint_entail}")
+
+                #2.(2):doc_corr
+                if joint_doc_entail == -1:
+                    if method=='ALCE':
+                        joint_doc_entail=joint_entail
+                    elif method=='baseline':
+                        joint_doc_entail=_run_nli_autoais(joint_passage, target_sent)
+                entail_doc+=joint_doc_entail
+                autoais_log[-1]["claim"]["doc_recall"] = joint_doc_entail
+                #print(f"the total time for two recall is {time.time() - start_time}")
+
+
+
+                # 4. 只有quote_corr=1（当该条数据，所有引用的拼接可以entail模型output的时候，）才计算引用粒度granularity
+                start_time=time.time()
+                if joint_entail:
+                    all_clauses = []
+                    clauses_first_three = []
+                    # 遍历每个不同的this_span
+                    #logger.info("calculating granularity")
+                    if len(ref_span)>5:
+                        print("Too many quotations!")
+                        autoais_log[-1]["claim"]["granularity_score"] = None
+                        autoais_log[-1]["claim"]["granularity_span"] = 0
+                    else:
+                        for idx, this_span in enumerate(ref_span):
+                            #logger.info(f"this span is {this_span}")
+                            # 分割引用跨度为子句
+                            clauses = re.split(r'([,.])', this_span)
+                            clauses = [clause.strip() for clause in clauses if
+                                       clause.strip() and any(char.isalnum() for char in clause.strip())]
+                            all_clauses.append(clauses)
+                            if idx<3:
+                                clauses_first_three.append(clauses)
+
+                        max_span_count = len(sum(all_clauses, []))
+                        if max_span_count==0:
+                            continue
+                        doc_count = len(all_clauses)
+                        min_comb_length=float('inf')
+
+                        if method=="ALCE" and grain=="default":
+                            gold_span_res=run_converge_2(clauses_first_three,target_sent)
+                        else:
+                            gold_span_res = run_converge_2(all_clauses, target_sent)
+                        # gold结果
+                        merged_gold_span_res = []
+
+                        # 遍历嵌套列表，并将其中的子列表合并到大列表中
+                        for sublist in gold_span_res:
+                            merged_gold_span_res.extend(sublist)
+                        autoais_log[-1]["claim"]["granularity_span"] = merged_gold_span_res
+                        min_comb_length=len(merged_gold_span_res)
+                        if min_comb_length!=float('inf'):
+                            granularity_score = min_comb_length / max_span_count
+                            granularity_list.append(granularity_score)
+                            autoais_log[-1]["claim"]["granularity_score"] = granularity_score
+
+
+                print(autoais_log[-1]["claim"]["granularity_span"])
+                print(autoais_log[-1]["claim"]["granularity_score"])
+                print(f"the total time for granularity is {time.time() - start_time}")
+            else:#不需要引用或没有引用
+                autoais_log[-1]['claim']['recall']=None
+                autoais_log[-1]["claim"]["granularity_score"]=None
+                autoais_log[-1]["claim"]["granularity_span"]=None
+
+
+        sent_total += len(sents)
+        ais_scores_need.append(correct_predictions / len(sents)) #是否正确判断需不需要引用：正确判断/总
+        if need_citations_sentences!=0: # recall：能entail的/需要引用的
+            ais_scores.append(entail / need_citations_sentences)
+            ais_doc_scores.append(entail_doc / need_citations_sentences)
+
+        #过滤None
+        granularity_list = [value for value in granularity_list if value is not None]
+
+    #logger.info(f"skipped {skipped}")
+    #autoais_log.append(f"skipped {skipped}")
+    ##print(autoais_log)
+    # print(ais_scores_need,ais_doc_scores,ais_scores,granularity_list)
+    return {
+        "citation_correct_prediction": 100 * np.mean(ais_scores_need),
+        "citation_doc_rec":100 * np.mean(ais_doc_scores),
+        "citation_quote_rec": 100 * np.mean(ais_scores),
+        "citation_granularity": 100 * np.mean(granularity_list)
+    } #autoais_log
+
+def compute_qampari_f1(data, cot=False):
+    prec = []
+    rec = []
+    rec_top5 = []
+    f1 = []
+    f1_top5 = []
+
+    num_preds = []
+    for item in data:
+        if cot:
+            if ":" in item['output']:
+                o = ':'.join(item['output'].split(":")[1:])  # try to separate the COT part and the answer list part.
+            else:
+                o = ""
+        else:
+            o = item['output']
+        preds = [normalize_answer(x.strip()) for x in remove_citations(o).rstrip().rstrip(".").rstrip(",").split(",")]
+        preds = [p for p in preds if len(p) > 0]  # delete empty answers
+        #print(preds)
+        num_preds.append(len(preds))
+        answers = [[normalize_answer(x) for x in ans] for ans in item['answers']]
+        flat_answers = [item for sublist in answers for item in sublist]
+        #print(flat_answers)
+        prec.append(sum([p in flat_answers for p in preds]) / len(preds) if len(preds) > 0 else 0)
+        #print(prec)
+        rec.append(sum([any([x in preds for x in a]) for a in answers]) / len(answers))
+        rec_top5.append(min(5, sum([any([x in preds for x in a]) for a in answers])) / min(5, len(answers)))
+        if (prec[-1] + rec[-1]) == 0:
+            f1.append(0)
+        else:
+            f1.append(2 * prec[-1] * rec[-1] / (prec[-1] + rec[-1]))
+        if (prec[-1] + rec_top5[-1]) == 0:
+            f1_top5.append(0)
+        else:
+            f1_top5.append(2 * prec[-1] * rec_top5[-1] / (prec[-1] + rec_top5[-1]))
+
+    return {
+        "num_preds": np.mean(num_preds),
+        "qampari_prec": 100 * np.mean(prec),
+        "qampari_rec": 100 * np.mean(rec),
+        "qampari_rec_top5": 100 * np.mean(rec_top5),
+        "qampari_f1": 100 * np.mean(f1),
+        "qampari_f1_top5": 100 * np.mean(f1_top5),
+    }
+
+def compute_length(data):
+    return sum(len(item['output'].split(' '))for item in data)/(len(data))
+    
+
+if __name__ =='__main__':
+    #question = "Why did New York City try to ban food donations to the poor?"
+    #output = "New York City, under Mayor Michael Bloomberg's administration, tried to ban food donations to the poor mainly due to concerns about the nutritional content of the donated food. The city argued that it couldn't inspect donated food for its salt, fat, and fiber content, thereby making it hard to control the nutritional quality of the food served to its homeless population [1][2][3]. Critics of this policy, however, have claimed such an approach demonstrated excessive control over people's eating habits and lacked common sense [2]. Despite the ban, many organizations like the New York City Rescue Mission continued to serve needy citizens through food donations [5]."
+    #compute_qa(question,output,['',''])
+    pass
+
+
+
+class Evaluator():
+    autoais_model_load = False
+
+    eval_criteria = {'test_pr':test_compute_autoais,'cite_recall_precision':compute_autoais, 'pr':compute_autoais,'qa':compute_qa,'rouge': compute_rouge_l,'claims':compute_claims, 'qampari':compute_qampari_f1,'length':compute_length,'str_em':compute_str_em,'grained':compute_autoais_grained,'cite_recall_precision_llm':lambda data: compute_autoais(data=data,entail_function=_run_llm_autoais),'mauve':compute_mauve}
+    def __init__(self,criteria= None, pipeline = None, ais_model = None) -> None:
+        self.eval_criteria = Evaluator.eval_criteria
+        self.pipeline = pipeline
+        self.get_data = {}
+        self.ais_model = ais_model
+        global ais_LLM
+        ais_LLM = ais_model
+
+            
+
+    def set_eval(self, eval_c, **data_get_key):
+        if eval_c in self.get_data.keys():
+            print(f'Already set! {eval_c}')
+            return 
+        if eval_c in self.eval_criteria.keys():
+            self.get_data[eval_c] = data_get_key
+            if eval_c == 'cite_recall_precision':
+                global autoais_model, autoais_tokenizer
+                if not Evaluator.autoais_model_load:
+                    print('Initializing eval model for citation precision and recall...') 
+                    try:
+                        autoais_model = AutoModelForSeq2SeqLM.from_pretrained(AUTOAIS_MODEL, torch_dtype=torch.bfloat16, device_map="auto")
+                        autoais_tokenizer = AutoTokenizer.from_pretrained(AUTOAIS_MODEL, use_fast=False)
+                        
+                    except:
+                        print('Unable to load model from hub, trying to load from local path...')
+                        autoais_model = AutoModelForSeq2SeqLM.from_pretrained(AUTOAIS_MODEL_ABSOLUTE, torch_dtype=torch.bfloat16, device_map="auto")
+                        autoais_tokenizer = AutoTokenizer.from_pretrained(AUTOAIS_MODEL_ABSOLUTE, use_fast=False)
+                    Evaluator.autoais_model_load = True
+            if eval_c == 'qa':
+                global qa_pipeline
+                qa_pipeline = transformers.pipeline("question-answering", model=QA_MODEL)
+
+        else:
+            raise KeyError('eval_criteria unavailable')
+    
+    def new_eval(self, name, eval_func, **data_get_key):
+        self.eval_criteria[name] = eval_func
+        self.set_eval(name, **data_get_key)
+
+    def __call__(self,data_from_pipeline= None):
+        result = {}
+        
+        for criteria, get_data in self.get_data.items():
+            if not data_from_pipeline:
+                data_dict = {}
+                for k, v in get_data.items():
+                    if isinstance(v,str):
+                        if v == 'output':
+                            data_dict[k] = ' '.join(self.pipeline.output)
+                        elif v == 'doc_cache':
+                            data_dict[k] = self.pipeline.doc_cache
+                        else:
+                            data_dict[k] = self.pipeline.dataset[self.pipeline.data_index][v]
+                    else:
+                        data_dict[k] = v
+            else:
+                data_dict = data_from_pipeline
+
+            eval_func = self.eval_criteria[criteria]
+            data = [data_dict]
+            result[criteria] = eval_func(data)
+        return result
+        
+
+
+class DefaultEvaluator(Evaluator):
+    def __init__(self, args = None, criteria= None, pipeline = None) -> None:
+        super().__init__(criteria,pipeline)
+        if args:
+            if  hasattr(args,'str_em') and args.str_em:
+                self.set_eval('str_em',output = PIPELINE_OUTPUT, qa_pairs = 'qa_pairs')
+            if  hasattr(args,'pr') and args.pr:
+                self.set_eval('cite_recall_precision', output = PIPELINE_OUTPUT, docs = PIPELINE_DOC_CACHE, question = 'question')
+            if  hasattr(args,'mauve') and args.mauve:
+                self.set_eval('mauve', output = PIPELINE_OUTPUT, answer = 'answer' ,question = 'question')
+            if  hasattr(args,'rouge') and args.rouge:
+                if (hasattr(args, 'dataset') and 'qampari' not in args.dataset.lower()) or not hasattr(args, 'dataset'):
+                    self.set_eval('rouge', output = PIPELINE_OUTPUT, answer = 'answer')
+            if  hasattr(args,'qa') and args.qa:
+                if (hasattr(args, 'dataset') and 'asqa' in args.dataset.lower()) or not hasattr(args, 'dataset'):
+                    self.set_eval('qa',output = PIPELINE_OUTPUT, qa_pairs = 'qa_pairs')
+            if  hasattr(args,'claims') and args.claims:
+                if (hasattr(args, 'dataset') and 'eli5' in args.dataset.lower()) or not hasattr(args, 'dataset'):
+                    self.set_eval('claims',output = PIPELINE_OUTPUT, claims = 'claims')
+            if  hasattr(args,'qampari') and args.qampari:
+                if (hasattr(args, 'dataset') and 'qampari' in args.dataset.lower()) or not hasattr(args, 'dataset'):
+                    self.set_eval('qampari',output = PIPELINE_OUTPUT, answers = 'answers')
+            if  hasattr(args,'length') and args.length:
+                self.new_eval('length',lambda data: len(data[0]['output'].split(' ')), output = PIPELINE_OUTPUT)
+
+        elif criteria:
+            if 'cite_recall_precision' in criteria:
+                self.set_eval('cite_recall_precision', output = PIPELINE_OUTPUT, docs = PIPELINE_DOC_CACHE, question = 'question')
+            if  hasattr(args,'mauve') and args.mauve:
+                self.set_eval('mauve', output = PIPELINE_OUTPUT, answer = 'answer' ,question = 'question')
+            if 'rouge' in criteria:
+                self.set_eval('rouge', output = PIPELINE_OUTPUT, answer = 'answer')
+            if 'qa' in criteria:
+                self.set_eval('qa',output = PIPELINE_OUTPUT, qa_pairs = 'qa_pairs')
+            if 'str_em' in criteria:
+                self.set_eval('str_em',output = PIPELINE_OUTPUT, qa_pairs = 'qa_pairs')
+            if 'claims' in criteria:
+                self.set_eval('claims',output = PIPELINE_OUTPUT, claims = 'claims')
+            if 'qampari' in criteria:
+                self.set_eval('qampari',output = PIPELINE_OUTPUT, answers = 'answers')
+            if 'length' in criteria:
+                self.new_eval('length',lambda data: len(data[0]['output'].split(' ')), output = PIPELINE_OUTPUT)
+
+        else:
+            self.new_eval('length',lambda data: len(data[0]['output'].split(' ')), output = PIPELINE_OUTPUT)

citekit/pipeline/pipeline.py

@@ -0,0 +1,423 @@
+from citekit.cite_modules.LLM import LLM,Module
+from citekit.cite_modules.augment_model import AugmentCluster, AttributingModule, MODEL_TYPE_MAPPING
+from citekit.prompt.prompt import ALCEVanillaPrompt, DocPrompt
+import logging
+import json
+from tqdm import tqdm
+import traceback
+import copy
+from citekit.utils.utils import flatten_dict
+import csv
+
+
+
+def merge_str_dicts(dicts):
+    result = {}
+    for dictionary in dicts:
+        for key, value in dictionary.items():
+            if key in result:
+                result[key] += ' ' + value
+            else:
+                result[key] = value
+    return result
+
+PIPELINE_OUTPUT = 'output'
+PIPELINE_DOC_CACHE = 'doc_cache'
+
+class DocCache():
+    def __init__(self) -> None:
+        self.__docs = list()
+    
+    def __len__(self):
+        return len(self.__docs)
+
+    def __getitem__(self,index):
+        if index>=0 and index <len(self):
+            return self.__docs[index]
+        else:
+            return None
+    
+    def get_last(self):
+        if self.__docs:
+            return self.__docs[-1]
+        
+    def add_doc(self, doc, add_id = True) -> int:
+        if not isinstance(doc, str):
+            assert isinstance(doc, dict) and 'text' in doc and 'title' in doc
+            doc = f'(Title: {doc["title"]}){doc["text"]}'
+        if add_id:
+            doc_head = f'Document [{len(self)+1}]'
+        else:
+            doc_head = ''
+        self.__docs.append(doc_head + doc)
+        return len(self)
+
+    def load_docs(self, docs, add_id = False):
+        for doc in docs:
+            self.add_doc(doc, add_id)
+        return len(self)
+    
+    def clear(self):
+        self.__docs = list()
+
+    def show_docs(self):
+        return self.__docs
+
+
+class Pipeline():
+    def __init__(self,save_path = None, sequence = None, head_prompt_maker = None, llm = None, module= None, retriever = None, evaluator = None, dataset = None, rich_eval = False, train_data = False, attributer = None) -> None:
+        self.save_path = save_path
+        self.train_data = train_data
+        self.head_prompt_maker = head_prompt_maker
+        self.table_head = True
+        self.attributer = attributer
+        self.llm = llm
+        self.initial_docs = None
+        self.data_keys = None
+        self.stored_clusters = []
+        self.module = []
+        if llm:
+            llm.connect_to(self)
+        if not isinstance(module,list) and module is not None:
+            if module:
+                module.connect_to(self)
+        else:
+            if isinstance(module, list):
+                for i in module:
+                    if isinstance(i, AugmentCluster) or isinstance(i, Module):
+                        i.connect_to(self)
+        self.dataset = dataset
+
+        self.data_index = 0
+        self.retriever = retriever
+        if retriever:
+            retriever.pipeline = self
+
+        self.eval = evaluator
+        if evaluator:
+            evaluator.pipeline = self
+        self.output = []
+        self.log = []
+        self.doc_cache = DocCache()
+        self.head = {}
+        self.result = {}
+        self.rich_eval = rich_eval
+        self.initial_module = None
+        
+    def load_data(self, dataset):
+        self.data = dataset
+
+    def set_initial_module(self, module):
+        self.initial_module = module
+    def get_initial_module(self):
+        return self.initial_module
+    
+    def set_data_keys(self, keys):
+        self.data_keys = keys
+    def get_data_keys(self):
+        return self.data_keys
+    
+
+    def update(self, update_object, config, update_info):
+        print(f'Updating {update_object} with {config} and {update_info}')
+        module = self.get_module_by_name(update_object)
+        if config in ['prompt', 'header']:
+            module.update(config, update_info)
+        elif config in ['destination']:
+            module.update(config, [self.get_module_by_name(update_info[0]), update_info[1]])
+        elif config in ['delete_destination']:
+            module.update(config, self.get_module_by_name(update_info))
+        elif config in ['new_model']:
+            model_type, model, key = update_info
+            print('Creating new model:', model_type, model)
+            new_model_class = MODEL_TYPE_MAPPING[model_type]
+            print('New model class:', new_model_class)
+            new_model = new_model_class(model)
+            new_model.connect_to(self)
+            print('Created new model:', new_model)
+            module.update('destination', [new_model, key])
+
+        else:
+            raise NotImplementedError
+
+
+
+    def set_initial_docs(self, d):
+        self.initial_docs = d
+    def get_initial_docs(self):
+        return self.initial_docs
+    
+    def run_on_dataset(self,datakeys,init_docs=None,initial_module= None,start=0):
+        if self.initial_module and not initial_module:
+            initial_module = self.initial_module
+        if self.save_path:
+            for i in range(start,len((self.dataset))):
+                self.data_index = i
+                try:
+                    self.run(datakeys,init_docs,initial_module,train=self.train_data)
+                except Exception as e:
+                    print(f'Error: {e}, skipping data {i}')
+                    traceback.print_exc()
+        else:
+            for i in range(start,len((self.dataset))):
+                self.data_index = i
+                try:
+                    self.run(datakeys,init_docs,initial_module,write=False,train=self.train_data)
+                except Exception as e:
+                    print(f'Error: {e}, skipping data {i}')
+                    traceback.print_exc()
+        
+            
+
+    def form_eval_data(self) -> dict:
+        """To write rich eval, you can use data from:
+        pipeline.dataset, doc_cache and output 
+        to post_process data as a argument dict for evaluation
+        """
+        raise NotImplementedError('You have to write <form_eval_data function> to apply rich eval with designed arguments.')
+
+    def direct_run(self, dynamic_prompt= {}, module = None):
+        if not module:
+            module = self.llm
+        if isinstance(module, AugmentCluster):
+            module = module.get_first_module()
+        while isinstance(module, Module):
+            if isinstance(dynamic_prompt,dict):
+                module.change_to_multi_process(False)
+                dynamic_prompt = module.generate(self.head,dynamic_prompt=dynamic_prompt)
+            elif isinstance(dynamic_prompt,list) and all([isinstance(d,dict) for d in dynamic_prompt]):
+                module.change_to_multi_process(True)
+                if module.parallel:
+                    dynamic_prompt = [module.generate(self.head,d) for d in dynamic_prompt]
+                    if module.merge:
+                        dynamic_prompt = merge_str_dicts(dynamic_prompt)
+                    module.add_output_to_head(module.last_message)
+                elif not module.iterative and not module.merge:
+                    for d in dynamic_prompt:
+                        self.direct_run(dynamic_prompt = d, module = copy.copy(module))
+                    #dynamic_prompt = [module.generate(self.head,d) for d in dynamic_prompt]
+                    break
+                elif module.iterative:
+                    iter_dynamic = {}
+                    for d in dynamic_prompt:
+                        iter_dynamic = module.generate(self.head,{**d,**iter_dynamic})
+                    dynamic_prompt = iter_dynamic
+                module.end_multi()
+            else:
+                print(type(dynamic_prompt))
+                raise TypeError(str(dynamic_prompt))
+            self.log.append(f'{module} -> {module.send()}\n: {module.last_message}')
+            if isinstance(module, Module):
+                module.output()
+                print('DEBUG:', str(module), module.end, module.turns, module.max_turn)
+                if module.end or module.turns > module.max_turn:
+                    break
+            module = module.send()
+            if isinstance(module, AugmentCluster):
+                module = module.get_first_module()
+
+
+    def __call__(self, data):
+        # run only one data
+        # backup
+        dataset_backup = self.dataset
+        current_data_index_backup = self.data_index
+        if hasattr(self,'current_data'):
+            current_data_backup = self.current_data
+        else:
+            current_data_backup = None
+
+        # set data and run
+        dataset = [data]
+        self.dataset = dataset
+        self.data_index = 0
+        result = self.run(datakeys = self.data_keys, init_docs = self.initial_docs, initial_module = self.initial_module, write = False, train = False)
+
+        # restore
+        self.data_index = current_data_index_backup
+        self.current_data = current_data_backup
+        self.dataset = dataset_backup
+
+
+        return result
+
+    def run(self, datakeys, init_docs = None, initial_module = None, write = True, train = False):
+        
+        # get data
+        self.current_data = self.dataset[self.data_index]
+        data = self.current_data
+
+        # from head prompt from specific data
+        head = dict()
+        for key in datakeys:
+            if isinstance(data[key],str):
+                head[key] = data[key]
+            else:
+                assert isinstance(data[key],list)
+                assert all([isinstance(item, str) for item in data[key]])
+                head[key] = ''.join(data[key])
+
+        #init
+        self.head = head
+        self.output = []
+        self.doc_cache.clear()
+        if init_docs:
+            self.doc_cache.load_docs(data[init_docs])
+        self.llm.reset()
+        if self.module:
+            for i in self.module:
+                i.reset()
+        self.log = []
+        # run only one data, and add data_index by 1
+        dynamic_prompt = {}
+        if not initial_module:
+            module = self.llm
+        else:
+            module = initial_module
+        if isinstance(module, AugmentCluster):
+            module = module.get_first_module()
+        while isinstance(module, Module):
+            if isinstance(dynamic_prompt,dict):
+                module.change_to_multi_process(False)
+                dynamic_prompt = module.generate(self.head,dynamic_prompt=dynamic_prompt)
+            elif isinstance(dynamic_prompt,list) and all([isinstance(d,dict) for d in dynamic_prompt]):
+                module.change_to_multi_process(True)
+                if module.parallel:
+                    dynamic_prompt = [module.generate(self.head,d) for d in dynamic_prompt]
+                    if module.merge:
+                        dynamic_prompt = merge_str_dicts(dynamic_prompt)
+                    module.add_output_to_head(module.last_message)
+                elif not module.iterative and not module.merge:
+                    for d in dynamic_prompt:
+                        self.direct_run(dynamic_prompt = d, module = copy.copy(module))
+                    #dynamic_prompt = [module.generate(self.head,d) for d in dynamic_prompt]
+                    break
+
+                elif module.iterative:
+                    iter_dynamic = {}
+                    for d in dynamic_prompt:
+                        iter_dynamic = module.generate(self.head,{**d,**iter_dynamic})
+                    dynamic_prompt = iter_dynamic
+                module.end_multi()
+            else:
+                print(type(dynamic_prompt))
+                raise TypeError(str(dynamic_prompt))
+            self.log.append(f'{module} -> {module.send()}\n: {module.last_message}')
+            if isinstance(module, Module):
+                module.output()
+                if module.end or module.turns > module.max_turn:
+                    break
+            module = module.send()
+            if isinstance(module, AugmentCluster):
+                module = module.get_first_module()
+
+        # if eval, send to evaluation
+        if self.eval: 
+            if not self.rich_eval:
+                self.result = self.eval()
+            else:
+                self.result = self.eval(self.form_eval_data())
+        else:
+            self.result = {}
+        if write:
+            self.write()
+        if train:
+            self.export_training_data()
+        
+        #self.logs = self.delete_inner_cluster_logs(self.log)
+        res = {'data':self.get_data(), 'doc_cache':self.doc_cache.show_docs(), 'log': self.log.copy(),'output':self.output,'result': self.result}
+        if self.attributer:
+            self.attributer.attribute_for_result(res)
+        self.data_index += 1
+        return res
+    
+    def delete_inner_cluster_logs(self, logs):
+        print(logs)
+        for cluster in self.stored_clusters:
+            cluster_name = str(cluster)
+            print('Combining logs for cluster:', cluster_name)
+            in_cluster = False
+            for i, log in enumerate(logs):
+                in_out_names = log.split('\n')[0]
+                if in_out_names in cluster_name:
+                    # This is the inner log
+                    if not in_cluster:
+                        in_cluster = True
+                        log_start = i
+                    else:
+                        continue
+                elif in_cluster:
+                    # This is the outer log
+                    in_cluster = False
+                    log_end = i
+                    cluster_output = logs[log_end]
+                    next_module = in_out_names.split('->')[1].strip()
+                    cluster_log = f"{cluster_name} -> {next_module}\n: {cluster_output}"
+                    logs = logs[:log_start] + [cluster_log] + logs[log_end+1:]
+        print('Final logs:', logs)
+        return logs
+        
+                    
+
+
+    def get_data(self):
+         return self.dataset[self.data_index]
+
+    def write(self):
+        '''Default writing'''
+        llm_token_used = self.llm.token_used
+        write_down = {'data':self.get_data(), 'doc_cache':self.doc_cache.show_docs(), 'log': self.log.copy(),'output':self.output,'result': self.result,'token_used':llm_token_used}
+
+        if self.attributer:
+            self.attributer.attribute_for_result(write_down)
+
+        with open(self.save_path, 'a', encoding='utf-8') as file:
+            json_line = json.dumps(write_down, indent=4)
+            file.write(json_line + '\n')
+
+    def get_module_by_name(self, name):
+        print('Getting module by name:', name)
+        for module in self.module:
+            if str(module) == name:
+                return module
+        if str(self.llm) == name:
+            return self.llm
+        
+        for cluster in self.stored_clusters:
+            print('trying cluster:', cluster)
+            if str(cluster) == name:
+                print('found cluster:', cluster)
+                return cluster
+    
+        return None
+    
+    def export_training_data(self):
+        flattened_data = [flatten_dict(self.result)]
+        header = set()
+        for item in flattened_data:
+            header.update(item.keys())
+        header = sorted(header)
+        with open('output.csv', mode='a', newline='') as file:
+            writer = csv.DictWriter(file, fieldnames = header)
+            if self.table_head:
+                writer.writeheader()
+                self.table_head = False
+            
+            for row in flattened_data:
+                writer.writerow(row)
+
+
+    def __str__(self) -> str:
+        return 'pipeline output'
+
+class Sequence(Pipeline):
+    def __init__(self, save_path=None, sequence=None, head_prompt_maker=None, retriever=None, evaluator=None, dataset=None, rich_eval=False) -> None:
+        first_module = sequence[0]
+        other = sequence[1:]
+        super().__init__(save_path, sequence, head_prompt_maker, first_module, other, retriever, evaluator, dataset, rich_eval)
+        for i in range(len(sequence)-1):
+            module = sequence[i]
+            assert isinstance(module, Module) or isinstance(module,AugmentCluster)
+            module.set_target(sequence[i+1],post_processing=lambda x: {module.output_as: x})
+        sequence[-1].set_output()
+        

citekit/prompt/prompt.py

@@ -0,0 +1,294 @@
+import json
+
+
+truncate = lambda x, l:  x[:l]
+token_len = len
+
+def combine(*args):
+    if all([isinstance(arg,dict) for arg in args]):
+        if len(args) == 1: 
+            return args[0]
+        else:
+            combined = args[0].copy()
+            combined.update(combine(*args[1:]))
+            return combined
+
+default_get = lambda key :  lambda data: data[key]
+
+
+
+class Prompt:
+    components = {}
+    template = ""
+    truncate = lambda x, l:  x[:l]
+    UNABLE = 'prompt_unable'
+
+
+    def update(self, **kargs):
+        try: 
+            for key in kargs.keys():
+                if key == 'template':
+                    arg_template = kargs[key]
+
+                if key == 'components':
+                    arg_components = kargs[key]
+            test_prompt = Prompt(arg_template, arg_components)
+        except Exception as e:
+            print(e)
+            print('Update Rejected due to invalid template or components')
+            return
+        
+        self.template = arg_template
+        self.components = arg_components
+
+    def __init__(self,template='', components={}, max_token=8000) -> None:
+
+        '''
+        Args:
+            template: The way to order and organize each components, use <NAME> to represent a component, <C1><C2>...<Cn>.
+            components: The content of a component, use {NAME} to represent the placeholder of corresponding data
+            max_token: a list as long as components, representing the max number of tokens for each component, or a int representing the same max_token for all components
+        '''
+
+        # template
+        self.template = template
+
+        # components
+        if isinstance(components,dict):
+            for key in components.keys():
+                if f'<{str(key)}>' not in self.template:
+                    raise Exception('component name not in template!')
+            self.components = components
+
+        # max_token
+        self.max_token = {}
+        if isinstance(max_token,list):
+            if len(components)==len(max_token):
+                self.max_token = {att:val for (att,val) in zip(components.keys(),max_token)}
+            else:
+                raise Exception('max_token is not corresponding to components')
+        elif isinstance(max_token,int):
+            self.max_token_init = max_token
+            self.max_token = {att:max_token for att in components}
+        else:
+            raise TypeError('max_token should be int or list')
+        
+    def __repr__(self) -> str:
+        prompt = self.template
+        for key in self.components.keys():
+            prompt = prompt.replace(f'<{str(key)}>',self.components[key])
+        return prompt
+    
+    def __str__(self) -> str:
+        return repr(self)
+
+    def part_template(self,**kargs):
+        '''
+        Add components in to the prompt.
+        '''
+        for part in kargs.keys():
+            if f'<{str(part)}>' in self.template:
+                self.components[part] = kargs[part]
+            else:
+                raise Exception('component name not in template!')
+    
+    def __call__(self, *args,**kargs) -> str:
+        return self.make_prompt(*args, **kargs)
+
+
+    def __str__(self):
+
+        input = {}
+        for key in self.components.keys():
+            input[key] = self.components[key]
+        
+
+        return self.make_prompt(input)
+
+
+    def make_prompt(self,*args,**kargs) -> str:
+        '''
+        arg: a dictionary containing all contents to the placeholder of the prompt
+        kargs: use NAME=value to pass arguments
+        '''
+
+        if args:
+            args = combine(*args)
+            args = args.copy()
+            args.update(kargs)
+        else:
+            args = kargs
+        prompt = self.template
+
+        for key in self.components.keys():
+            if key not in args or args[key] == Prompt.UNABLE:
+                prompt = prompt.replace(f'<{str(key)}>',"")
+            else:
+                prompt = prompt.replace(f'<{str(key)}>', self.components[key])
+        
+        prompt_args = {}
+        for key in args.keys():
+            if key in self.components.keys():
+                if self.max_token.get(key):
+                    max_token = self.max_token.get(key)
+                else:
+                    max_token = min(4096,self.max_token_init)
+                if token_len(args[key])> max_token:
+                    args[key] = self.truncate(args[key],max_token)
+
+
+            return prompt.format(**args)
+    
+    def set_max_token(self,**kargs) -> None:
+        for key in kargs.keys():
+            if key in self.components.keys():
+                self.max_token[key] = kargs.get(key)
+            else:
+                raise KeyError(f'{key} not in Template!')
+            
+    def load_data(self,data_loader,*keys,**projections):
+        '''
+        load data to make prompts from a data loader
+        projections: the function to get the information from a data.
+        '''
+
+        prompts = []
+        for data in data_loader:
+            l_contents = {key:default_get(key)(data) for key in keys}
+            d_contents = {projection:projections[projection](data) for projection in projections.keys()}
+            prompts.append(self.make_prompt({**l_contents, **d_contents}))
+
+        return prompts
+    
+
+
+
+
+class DocPrompt(Prompt):
+    '''
+    Containing Doc ID, Title and Passage in order:
+    
+    Document:[{ID}]
+    (Title:{Title}) 
+    {Passage}
+    '''
+    def __init__(self, template='<ID><Title><Passage>', components={'ID':'Document[{ID}]: ','Title':'(Title:{Title})','Passage':'{Passage}\n'}, max_token=4096) -> None:
+        super().__init__(template, components, max_token)
+
+
+class ALCEDocPrompt(Prompt):
+    '''
+    Containing Doc ID, Title and Passage in order:
+    
+    Document:[{ID}]
+    (Title:{Title}) 
+    {Passage}
+    '''
+    def __init__(self, template='<ID><title><text>', components={'ID':'Document [{ID}]','title':'(Title:{title}): ','text':'{text}\n'}, max_token=4096) -> None:
+        super().__init__(template, components, max_token)
+
+    def default_load_data(self,data_loader, text = 'text', from_idx = 0):
+        return super().load_data(list(enumerate(data_loader)),text = lambda data: data[1][text],ID = lambda data: str(data[0]+1 + from_idx),title = lambda data: data[1]['title'])
+
+    def default_load_data_wo_ID(self,data_loader):
+        return super().load_data(list(enumerate(data_loader)),text = lambda data: data[1]['text'],title = lambda data: data[1]['title'])
+    def default_load_data_wo_title(self,data_loader):
+        return super().load_data(list(enumerate(data_loader)),text = lambda data: data[1]['text'],ID = lambda data: str(data[0]+1))
+    def default_load_data_extraction(self,data_loader):
+        return super().load_data(list(enumerate(data_loader)),text = lambda data: data[1]['extraction'],ID = lambda data: str(data[0]+1),title = lambda data: data[1]['title'])
+    def default_load_data_summary(self,data_loader):
+        return super().load_data(list(enumerate(data_loader)),text = lambda data: data[1]['summary'],ID = lambda data: str(data[0]+1),title = lambda data: data[1]['title'])
+
+class ALCEVanillaPrompt(Prompt):
+    '''
+    Containing INST(Instruction), Question, Doc and Answer in order:
+
+    {INST}
+
+    Question:{Question}
+
+    {Doc}
+    Answer:{Answer}
+    '''
+    def __init__(self,
+                template="<INST><Question><Doc><Answer>\n", 
+                components={'INST':'{INST}\n\n', 'Question':'Question:{Question}\n\n','Doc':'{Doc}\n','Answer':'Answer:{Answer}'}, 
+                max_token=4096) -> None:
+        super().__init__(template, components, max_token)
+
+class NewALCEVanillaPrompt(Prompt):
+    '''
+    Containing INST(Instruction), Question, Doc and Answer in order:
+
+    {INST}
+
+    Question:{Question}
+
+    {Doc}
+    Answer:{Answer}
+    '''
+    def __init__(self,
+                template="<INST><question><docs><answer>\n", 
+                components={'INST':'{INST}\n\n', 'question':'Question:{question}\n\n','docs':'{docs}\n','answer':'Answer:{answer}'}, 
+                max_token=4096) -> None:
+        super().__init__(template, components, max_token)
+
+
+
+class AGEEPrompt(Prompt): 
+    '''
+    Containing INST(Instruction), Question and Doc in order:
+
+    {INST}
+
+    Question:{Question}
+
+    Search Results:{Doc}
+    '''
+    def __init__(self,
+                template="<INST><Question><Doc>\n", 
+                components={'INST':'{INST}\n\n', 'Question':'Question:\n{Question}\n','Doc':'Search Results:\n{Doc}\n'}, 
+                max_token=4096) -> None:
+        super().__init__(template, components, max_token)
+
+
+
+
+alce_prompt= ALCEVanillaPrompt()
+#alce_prompt.set_max_token(INST = 10,Doc = 100,Answer = 15)
+DocP= DocPrompt()
+
+
+#print(content['demos'])
+
+
+#print(data[0])
+'''
+pps = alce_prompt.load_data(content['demos'],
+            INST = lambda _: content['instruction'],
+            Question = lambda data: data['question'],
+            Doc = lambda data: ''.join(DocPrompt().load_data(list(enumerate(data['docs'])),
+                                                              ID = lambda data: str(data[0]),
+                                                              Title = lambda data: data[1]['title'],
+                                                              Passage = lambda data: data[1]['text'])),
+            Answer = lambda data: data['answer'])
+
+'''
+
+#print(pps[0])
+
+'''
+data_loader = []
+with open('data.txt','r',encoding='utf-8') as f:
+    content = f.readlines()
+    for i,c in enumerate(content):
+        if i%3 == 0:
+            data_loader.append({'Q':c.strip(),'A':content[i+1].strip()})
+print(data_loader)
+
+
+pps = Dp.load_data(data_loader,
+             INST= lambda data: "Instruction: Write an accurate, engaging, and concise answer for the given question",
+             Question= lambda data: data['Q'],
+             Answer= lambda data: data['A'])
+'''

citekit/utils/utils.py

@@ -0,0 +1,317 @@
+import numpy as np
+import string
+import re
+import collections
+import torch
+import nltk
+
+def one_paragraph(text):
+    paras = text.lstrip('\n').split('\n\n')
+    if not paras:
+        return ''
+    else:
+        return paras[0].rstrip('\n')
+
+def strong_one_paragraph(text):
+    paras = text.lstrip('\n').split('\n')
+    if not paras:
+        return ''
+    else:
+        return paras[0].rstrip('\n')
+    
+def compute_str_em(data):
+    """Compute STR-EM metric (only for ASQA)
+    Args:
+        data: requires field `qa_pairs/short_answers` and `output`
+    Returns:
+        STR-EM and STR-EM-HIT ()
+    """
+    if 'qa_pairs' not in data[0] or data[0]['qa_pairs'] is None:
+        return 0
+
+    acc = []
+    hit = []
+
+    for item in data:
+        loc_acc = []
+        for qa_pair in item['qa_pairs']:
+            loc_acc.append(exact_presence(qa_pair['short_answers'], item["output"]))
+        acc.append(np.mean(loc_acc))
+        hit.append(int(np.mean(loc_acc) == 1))
+
+    return 100 * np.mean(acc)
+    return 100 * np.mean(acc), 100 * np.mean(hit)
+
+def average(func):
+    def avg_func(dataset):
+        print(len(dataset))
+        results = [func(*data) for data in dataset] if dataset else []
+        if results:
+            return np.mean(np.array(results), axis=0).tolist()
+        else:
+            return None
+    return avg_func
+
+def normalize_answer(s):
+    def remove_articles(text):
+        return re.sub(r"\b(a|an|the)\b", " ", text)
+
+    def white_space_fix(text):
+        return " ".join(text.split())
+
+    def remove_punc(text):
+        exclude = set(string.punctuation)
+        return "".join(ch for ch in text if ch not in exclude)
+
+    def lower(text):
+        return text.lower()
+
+    return white_space_fix(remove_articles(remove_punc(lower(s))))
+
+def compute_f1(a_gold, a_pred):
+    """Compute F1 score between two strings."""
+
+    def _get_tokens(s):
+        if not s:
+            return []
+        return normalize_answer(s).split()
+
+    gold_toks = _get_tokens(a_gold)
+    pred_toks = _get_tokens(a_pred)
+
+    common = collections.Counter(gold_toks) & collections.Counter(pred_toks)
+    num_same = sum(common.values())
+
+    if len(gold_toks) == 0 or len(pred_toks) == 0:
+        # If either is no-answer, then F1 is 1 if they agree, 0 otherwise
+        return int(gold_toks == pred_toks)
+
+    if num_same == 0:
+        return 0
+
+    precision = 1.0 * num_same / len(pred_toks)
+    recall = 1.0 * num_same / len(gold_toks)
+    f1 = (2 * precision * recall) / (precision + recall)
+
+    return f1
+
+
+def compute_exact(a_gold, a_pred):
+    """Check whether two strings are equal up to normalization."""
+
+    return int(normalize_answer(a_gold) == normalize_answer(a_pred))
+
+
+def exact_presence(short_answers, context):
+    """Verify if any of the answers is present in the given context.
+    Args:
+        short_answers: list of short answers to look for in the context
+        context: a paragraph to search for short answers
+    Returns:
+        true if any of the short answers is present in the context
+    """
+
+    n_short_answers = [normalize_answer(sa) for sa in short_answers]
+    n_context = normalize_answer(context)
+
+    for ans in n_short_answers:
+        if ans in n_context:
+            return True
+
+    return False
+
+def output_begin_with(word):
+    def f(self) -> bool:
+        return self.last_message.strip().lower()[:len(word)] == word
+    return f
+
+def output_end_with(word):
+    def f(self) -> bool:
+        return strong_one_paragraph(self.last_message.strip())[-len(word):] == word
+    return f
+
+
+def make_as(datakey):
+    def f(passage):
+        return {datakey:passage}
+    return f
+
+def cut_and_make_as(datakey):
+    def f(passage):
+        return {datakey:one_paragraph(passage)}
+    return f
+
+def remove_citations(sent):
+    return re.sub(r"{\d+", "", re.sub(r" {\d+", "", sent)).replace(" |", "").replace("}", "").replace("{", "")
+
+def remove_citations(sent):
+    return re.sub(r"\[\d+", "", re.sub(r" \[\d+", "", sent)).replace(" |", "").replace("]", "")
+
+
+def match_document(ref_mark, output_ref_span):
+    ref = set()
+    ref_span = []
+    for num in ref_mark:
+        ref_str = str(num)
+        if ref_str in output_ref_span:
+            ref_parts = output_ref_span[ref_str].split("[")
+            if len(ref_parts) > 1:
+                ref_id_parts = ref_parts[1].split("]")
+                if len(ref_id_parts) > 0:
+                    ref_id = ref_id_parts[0].strip()
+                    if ref_id.isdigit():
+                        ref.add(int(ref_id))  # 添加Document id
+
+            ref_span_parts = output_ref_span[ref_str].split(":",1)#第一个冒号后面的片段
+            if len(ref_span_parts) > 1:
+                ref_span.append(ref_span_parts[1].strip())  # 添加后面的句子片段
+            else:
+                ref_span.append('')
+    return list(ref), ref_span
+
+def get_max_memory():
+    """Get the maximum memory available for the current GPU for loading models."""
+    free_in_GB = int(torch.cuda.mem_get_info()[0]/1024**3)
+    max_memory = f'{free_in_GB-6}GB'
+    n_gpus = torch.cuda.device_count()
+    max_memory = {i: max_memory for i in range(n_gpus)}
+    return max_memory
+
+def each_make_as(key):
+    def function(output):
+        sents = nltk.sent_tokenize(one_paragraph(output))
+        if len(sents)>3:
+            sents = sents[:3]
+        return [make_as(key)(sent) for sent in sents]
+    return function
+
+def each_par_make_as(key):
+    def function(output):
+        sents = one_paragraph(output).split('\n')
+        if len(sents)>3:
+            sents = sents[:3]
+        return [make_as(key)(sent) for sent in sents]
+    return function
+
+def sentence(key):
+    def function(output):
+        sents = nltk.sent_tokenize(one_paragraph(output))
+        for sent in sents:
+            refs = re.findall(r'\[\d+\]', sent)
+            if refs:
+                return make_as(key)(sent)
+        return make_as(key)('')
+    return function
+
+def sentences(key):
+    def function(output):
+        sents = nltk.sent_tokenize(one_paragraph(output))
+        return [make_as(key)(sent) for sent in sents][:1]
+    return function
+
+def three_sentences(key):
+    def function(output):
+        sents = nltk.sent_tokenize(one_paragraph(output))
+        return [make_as(key)(sent) for sent in sents][:3]
+    return function
+
+def first_sentence(text):
+    sents = nltk.sent_tokenize(one_paragraph(text))
+    for sent in sents:
+        return sent
+    return ''
+
+def flatten_dict(d, parent_key='', sep='_'):
+    items = []
+    for k, v in d.items():
+        new_key = f'{parent_key}{sep}{k}' if parent_key else k
+        if isinstance(v, dict):
+            items.extend(flatten_dict(v, new_key, sep=sep).items())
+        else:
+            items.append((new_key, v))
+    return dict(items)
+
+
+
+import re
+from bs4 import BeautifulSoup
+
+def parse_html_prompt(input_str):
+    soup = BeautifulSoup(input_str, "html.parser")
+    
+    # 处理 <p></p> 内的内容
+    p_content = soup.find("p").decode_contents().replace("<br>", "\n")
+    p_content = re.sub(r'<span[^>]*>(.*?)</span>', r'<\1>', p_content)
+    template = p_content.strip().replace(' <br/>', '').replace(' ', '').replace('<br/>', '')
+    
+    # 解析 component-item
+    components = {}
+    for item in soup.find_all("div", class_="component-item"):
+        key_span = item.find("div", class_="component-key").find("span")
+        key = key_span.get_text(strip=True) if key_span else ""
+        value_div = item.find("div", class_="component-value")
+        value_content = value_div.decode_contents()
+        value_content = re.sub(r'<span[^>]*>(.*?)</span>', r'{\1}', value_content)
+        components[key] = value_content.strip().replace(' <br/>', '').replace('<br/>', '')
+    
+    # 解析 self-info-item
+    self_prompt = {}
+    for item in soup.find_all("div", class_="self-info-item"):
+        key_span = item.find("div", class_="component-key").find("span")
+        key = key_span.get_text(strip=True) if key_span else ""
+        value_div = item.find("div", class_="component-value")
+        value = value_div.get_text(strip=True) if value_div else ""
+        self_prompt[key] = value.replace(' <br/>', '').replace('<br/>', '')
+    
+    return {
+        'template': template,
+        'components': components,
+        'self_prompt': self_prompt
+    }
+
+
+def parse_html_destination(input_str):
+    soup = BeautifulSoup(input_str, "html.parser")
+    destination = soup.find("destination").get_text(strip=True)
+    prompt_key = soup.find("prompt_key").get_text(strip=True)
+    return destination, prompt_key
+
+def parse_html_new_model(input_str):
+    soup = BeautifulSoup(input_str, "html.parser")
+    model_type = soup.find("model_type").get_text(strip=True)
+    model_name = soup.find("model").get_text(strip=True)
+    key = soup.find("prompt_key").get_text(strip=True)
+    return model_type, model_name, key
+
+def parse_delete_destination(input_str):
+    soup = BeautifulSoup(input_str, "html.parser")
+    destination = soup.find("deletedestination").get_text(strip=True)
+    return destination
+
+def parse_html_header(input_str):
+    soup = BeautifulSoup(input_str, "html.parser")
+    header = soup.find("to_head").get_text(strip=True)
+    return header
+
+def parse_html_config(info):
+    config = ''
+    if 'class="component-value"' in info:
+        func = parse_html_prompt
+        config = 'prompt'
+    elif '</destination>' in info:
+        func = parse_html_destination
+        config = 'destination'
+    elif '<model_type>' in info:
+        func = parse_html_new_model
+        config = 'new_model'
+    elif 'deletedestination' in info:
+        config = 'delete_destination'
+        func = parse_delete_destination
+    elif 'to_head' in info:
+        config = 'header'
+        func = parse_html_header
+    else:
+        raise NotImplementedError
+    result = func(info)
+    print(info, 'parsed as', config)
+    return config, result

context_cite/__init__.py

@@ -0,0 +1,4 @@
+from .context_citer import ContextCiter
+
+__version__ = "0.0.1"
+VERSION = __version__