TruthTorchLM 0.0.0__py3-none-any.whl

TruthTorchLM/__init__.py

@@ -0,0 +1,16 @@
+from .truth_methods.truth_method import TruthMethod
+from TruthTorchLM import utils ##TODO do we really need to import this?
+from TruthTorchLM import scoring_methods
+from TruthTorchLM import truth_methods
+from .generation import generate_with_truth_value
+from .calibration import calibrate_truth_method
+from TruthTorchLM import evaluators
+from .evaluators import evaluate_truth_method
+from .templates import DEFAULT_USER_PROMPT, DEFAULT_SYSTEM_PROMPT ##TODO import all?
+from .availability import AVAILABLE_DATASETS, AVAILABLE_EVALUATION_METRICS
+from TruthTorchLM import normalizers
+
+from TruthTorchLM import long_form_generation
+
+
+#__all__ = ['generate_with_truth_value']

TruthTorchLM/availability.py

@@ -0,0 +1,14 @@
+AVAILABLE_API_MODELS = ['gpt-4o', 'gpt-4o-2024-05-13', 'gpt-4o-2024-08-06', 'chatgpt-4o-latest', 'gpt-4o-mini', 'gpt-4o-mini-2024-07-18', 
+                        'gpt-4-turbo','gpt-4-turbo-2024-04-09', 'gpt-4-turbo-preview', 'gpt-4-0125-preview', 'gpt-4-1106-preview', 'gpt-4', 
+                        'gpt-4-0613', 'gpt-4-0314', 'gpt-3.5-turbo-0125', 'gpt-3.5-turbo', 'gpt-3.5-turbo-1106', 'gpt-3.5-turbo-instruct', 'together_ai/togethercomputer/llama-2-70b'] 
+
+PROB_AVAILABLE_API_MODELS = ['gpt-4o', 'gpt-4o-2024-05-13', 'gpt-4o-2024-08-06', 'chatgpt-4o-latest', 'gpt-4o-mini', 'gpt-4o-mini-2024-07-18', 
+                        'gpt-4-turbo','gpt-4-turbo-2024-04-09', 'gpt-4-turbo-preview', 'gpt-4-0125-preview', 'gpt-4-1106-preview', 'gpt-4', 
+                        'gpt-4-0613', 'gpt-4-0314', 'gpt-3.5-turbo-0125', 'gpt-3.5-turbo', 'gpt-3.5-turbo-1106', 'gpt-3.5-turbo-instruct', 'together_ai/togethercomputer/llama-2-70b'] 
+
+ACTIVATION_AVAILABLE_API_MODELS = []
+
+AVAILABLE_DATASETS = ['trivia_qa', 'gsm8k', 'natural_qa', 'pop_qa', 'simple_qa']
+LONG_FORM_AVAILABLE_DATASETS = ['longfact_concepts', 'longfact_objects']
+
+AVAILABLE_EVALUATION_METRICS = ['auroc', 'auprc', 'auarc', 'accuracy', 'f1', 'precision', 'recall', 'prr']

TruthTorchLM/calibration.py

@@ -0,0 +1,36 @@
+from transformers import PreTrainedModel, PreTrainedTokenizer, PreTrainedTokenizerFast
+from typing import Union
+from TruthTorchLM.truth_methods import TruthMethod
+from TruthTorchLM.evaluators import CorrectnessEvaluator, ROUGE
+from TruthTorchLM.templates import DEFAULT_SYSTEM_BENCHMARK_PROMPT, DEFAULT_USER_PROMPT 
+from TruthTorchLM.utils.dataset_utils import get_dataset
+from TruthTorchLM.utils.eval_utils import run_over_dataset
+import numpy as np
+
+
+
+def calibrate_truth_method(dataset: Union[str, list], model:Union[str,PreTrainedModel],  truth_methods: list[TruthMethod], tokenizer: Union[PreTrainedTokenizer, PreTrainedTokenizerFast] =None,
+correctness_evaluator:CorrectnessEvaluator = ROUGE(0.7), size_of_data:float = 1.0, previous_context:list =[{'role': 'system', 'content': DEFAULT_SYSTEM_BENCHMARK_PROMPT}], 
+ user_prompt:str = DEFAULT_USER_PROMPT, seed:int = 0, wandb_run = None, return_method_details:bool = False, wandb_push_method_details:bool = False, split = 'train', **kwargs):
+    
+    dataset = get_dataset(dataset, size_of_data=size_of_data, seed=seed, split = split)
+
+    output_dict = run_over_dataset(dataset, model, truth_methods, tokenizer = tokenizer, correctness_evaluator = correctness_evaluator, 
+                                   previous_context = previous_context, user_prompt = user_prompt, seed = seed, return_method_details = return_method_details, 
+                                   wandb_run = wandb_run, wandb_push_method_details = wandb_push_method_details, **kwargs)
+
+    for i, truth_method in enumerate(truth_methods):
+        truth_values = output_dict[f'truth_method_{i}']['truth_values']
+        truth_values = np.array(truth_values)
+        truth_values[np.isnan(truth_values)] = 0
+        correctness = output_dict['generation_correctness']
+        #if generation_correctness is -1, it means that the model didn't attempt to generate an answer, remove those from the evaluation
+        truth_method.normalizer.calibrate(generation_performance_scores=correctness, truth_values=truth_values)
+    return output_dict 
+
+
+    
+
+
+
+    

TruthTorchLM/evaluators/__init__.py

@@ -0,0 +1,8 @@
+from .correctness_evaluator import CorrectnessEvaluator
+from .rouge import ROUGE
+from .bleu import BLEU
+from .model_judge import ModelJudge
+from .eval_truth_method import evaluate_truth_method, get_metric_scores
+
+__all__ = ['CorrectnessEvaluator', 'ROUGE', 'BLEU', 'evaluate_truth_method', 'ModelJudge', 'get_metric_scores']
+

TruthTorchLM/evaluators/bleu.py

@@ -0,0 +1,20 @@
+
+from .correctness_evaluator import CorrectnessEvaluator
+import evaluate
+
+
+class BLEU(CorrectnessEvaluator):
+    def __init__(self, threshold: float = 0.5): 
+        super().__init__()
+        self.threshold = threshold  
+        self.bleu = evaluate.load('bleu')
+
+    def __call__(self, question_text:str, generated_text: str,  ground_truths: list[str], seed:int = None) -> bool:
+        for i in range(len(ground_truths)):
+            bleu_results = self.bleu.compute(predictions = [generated_text], references=[ground_truths[i]])
+            if bleu_results['bleu'] > self.threshold:
+                return 1
+        return 0
+    
+    def __str__(self):
+        return f"BLEU with threshold {self.threshold}"

TruthTorchLM/evaluators/correctness_evaluator.py

@@ -0,0 +1,14 @@
+from abc import ABC, abstractmethod
+
+class CorrectnessEvaluator(ABC):
+    
+    def __init__(self):
+        pass
+
+    @abstractmethod
+    def __call__(self, question_text:str, generated_text: str,  ground_truth_text: list[str], seed:int = None) -> int:
+        raise NotImplementedError("Subclasses must implement this method")
+    
+    @abstractmethod
+    def __str__(self):
+        raise NotImplementedError("Subclasses must implement this method")

TruthTorchLM/evaluators/eval_truth_method.py

@@ -0,0 +1,59 @@
+from transformers import PreTrainedModel, PreTrainedTokenizer, PreTrainedTokenizerFast
+from typing import Union
+from TruthTorchLM.truth_methods import TruthMethod
+from .correctness_evaluator import CorrectnessEvaluator 
+from .rouge import ROUGE
+from TruthTorchLM.availability import AVAILABLE_EVALUATION_METRICS
+from TruthTorchLM.templates import DEFAULT_SYSTEM_BENCHMARK_PROMPT, DEFAULT_USER_PROMPT
+from TruthTorchLM.utils.dataset_utils import get_dataset
+from TruthTorchLM.utils.eval_utils import metric_score, run_over_dataset
+import wandb
+
+
+def evaluate_truth_method(dataset: Union[str, list], model:Union[str,PreTrainedModel],  truth_methods: list[TruthMethod], tokenizer: Union[PreTrainedTokenizer, PreTrainedTokenizerFast]=None, eval_metrics:list[str] = ['auroc'],
+                          correctness_evaluator:CorrectnessEvaluator = ROUGE(0.7), size_of_data = 1.0,  previous_context:list =[{'role': 'system', 'content': DEFAULT_SYSTEM_BENCHMARK_PROMPT}], 
+                          user_prompt:str = DEFAULT_USER_PROMPT, seed:int = 0, return_method_details:bool = False, wandb_run = None, wandb_push_method_details:bool = False, 
+                          batch_generation=True,  add_generation_prompt = True, continue_final_message = False, split='test',  **kwargs):
+    
+    dataset = get_dataset(dataset, size_of_data=size_of_data, seed=seed, split = split)
+
+    for eval_metric in eval_metrics:
+        if eval_metric not in AVAILABLE_EVALUATION_METRICS:
+            raise ValueError(f"Evaluation metric {eval_metric} is not available. Available evaluation metrics are: {AVAILABLE_EVALUATION_METRICS}")
+                
+    output_dict = run_over_dataset(dataset, model, truth_methods, tokenizer = tokenizer, correctness_evaluator = correctness_evaluator, 
+                                   previous_context = previous_context, user_prompt = user_prompt, seed = seed, return_method_details = return_method_details, 
+                                   wandb_run = wandb_run, wandb_push_method_details= wandb_push_method_details, 
+                                   batch_generation=batch_generation, add_generation_prompt=add_generation_prompt, continue_final_message=continue_final_message, **kwargs)
+
+    eval_list = get_metric_scores(output_dict=output_dict, eval_metrics=eval_metrics, seed=seed)
+
+    
+    if wandb_run:
+        wandb_run.log({'model_accuracy': sum(output_dict['generation_correctness'])/len(output_dict['generation_correctness'])})
+        
+        eval_dict = eval_list[0]
+        for key, _ in eval_dict.items():
+            methods = []
+            scores = []
+            for i, cur_eval_dict in enumerate(eval_list):
+                score = cur_eval_dict[key]
+                scores.append(score)
+                methods.append(str(truth_methods[i].__class__.__name__))
+                wandb_run.log({f'{key}_of_method_{i}_{str(truth_methods[i].__class__.__name__)}': score})
+
+            data = [[method, score] for (method, score) in zip(methods, scores)]
+            table = wandb.Table(data=data, columns = ["methods", "scores"])
+            wandb.log({f"{key}" : wandb.plot.bar(table, "methods", "scores",
+                            title=f"{key} Scores of Truth Methods")})  
+                
+    return {'eval_list': eval_list, 'output_dict': output_dict}
+
+
+def get_metric_scores(output_dict:dict, eval_metrics:list[str], seed:int=0):
+    truth_methods = output_dict['truth_methods']
+    eval_list = []
+    for i in range(len(truth_methods)):
+        eval_dict = metric_score(eval_metrics, output_dict['generation_correctness'], output_dict[f'truth_method_{i}']['truth_values'], output_dict[f'truth_method_{i}']['normalized_truth_values'], seed=seed)
+        eval_list.append(eval_dict)
+    return eval_list

TruthTorchLM/evaluators/model_judge.py

@@ -0,0 +1,61 @@
+from .correctness_evaluator import CorrectnessEvaluator
+import evaluate
+from typing import Union
+from transformers import PreTrainedModel, PreTrainedTokenizer, PreTrainedTokenizerFast
+from litellm import completion
+import random
+import torch
+from TruthTorchLM.templates import DEFAULT_JUDGE_PROMPT, DEFAULT_JUDGE_SYSTEM_PROMPT
+
+
+class ModelJudge(CorrectnessEvaluator):
+    def __init__(self, model:Union[PreTrainedModel, str], tokenizer:Union[PreTrainedTokenizer, PreTrainedTokenizerFast] = None, prompt:str = DEFAULT_JUDGE_PROMPT, system_prompt:str = DEFAULT_JUDGE_SYSTEM_PROMPT, num_retries:int = 1) -> None: 
+        super().__init__()
+        self.model = model
+        self.tokenizer = tokenizer
+        self.prompt = prompt
+        self.system_prompt = system_prompt
+        self.num_retries = num_retries
+        
+
+    def __call__(self, question_text:str, generated_text: str,  ground_truths: list[str], seed:int = None) -> bool:
+        if seed == None:
+            seed = random.randint(0, 1000000)
+
+        chat = [{"role": "system", "content": self.system_prompt},
+        {"role": "user", "content": self.prompt.format(question = question_text, ground_truths = ', '.join(ground_truths), answer = generated_text)}]
+        if type(self.model) == str:
+            response = completion(
+                        model=self.model,
+                        messages=chat,
+                        seed=seed,
+                        num_retries=self.num_retries
+                    )
+            generated_text = response.choices[0].message['content']
+        else:
+            torch.manual_seed(seed)
+            random.seed(seed)
+            text = self.tokenizer.apply_chat_template(chat, tokenize = False)
+            input_ids = self.tokenizer.encode(text, return_tensors="pt").to(self.model.device)
+            model_output = self.model.generate(input_ids)
+            tokens = model_output[0][len(input_ids[0]):]
+            generated_text = self.tokenizer.decode(tokens, skip_special_tokens = False)
+
+        if 'incorrect' in generated_text.lower():
+            return 0
+        elif 'correct' in generated_text.lower(): 
+            return 1
+        elif "not_attempted" in generated_text.lower():
+            return -1
+        else:
+            #output warning
+            print("The output of the judge model is not in the expected format. Not attempted will be returned.")
+            return -1
+
+
+            
+            
+        
+    
+    def __str__(self):
+        return f"ROUGE with threshold {self.threshold} and type {self.rouge_type}"

TruthTorchLM/evaluators/rouge.py

@@ -0,0 +1,19 @@
+from .correctness_evaluator import CorrectnessEvaluator
+import evaluate
+
+class ROUGE(CorrectnessEvaluator):
+    def __init__(self, threshold: float = 0.5, rouge_type: str = 'rougeL'): 
+        super().__init__()
+        self.threshold = threshold  
+        self.rouge = evaluate.load('rouge')
+        self.rouge_type = rouge_type
+
+    def __call__(self, question_text:str, generated_text: str,  ground_truths: list[str], seed:int = None) -> bool:
+        for i in range(len(ground_truths)):
+            rouge_results = self.rouge.compute(predictions = [generated_text], references=[ground_truths[i]])
+            if rouge_results[self.rouge_type] > self.threshold:
+                return 1
+        return 0
+    
+    def __str__(self):
+        return f"ROUGE with threshold {self.threshold} and type {self.rouge_type}"

TruthTorchLM/generation.py

@@ -0,0 +1,389 @@
+
+import copy
+import torch
+import random
+from typing import Union
+from litellm import completion
+
+from transformers import PreTrainedModel, PreTrainedTokenizer, PreTrainedTokenizerFast
+
+#from .truth_methods.truth_method import TruthMethod
+from TruthTorchLM.availability import AVAILABLE_API_MODELS, PROB_AVAILABLE_API_MODELS
+from TruthTorchLM.utils.common_utils import generate, fix_tokenizer_chat
+
+import time
+
+
+def generate_with_truth_value(model:Union[PreTrainedModel, str], messages:list, question_context:str = None, truth_methods: list = [], tokenizer: Union[PreTrainedTokenizer, PreTrainedTokenizerFast]= None,  
+generation_seed=None, batch_generation=True, add_generation_prompt = True, continue_final_message = False,  **kwargs)-> dict:
+    if type(model) == str:
+        return generate_with_truth_value_api(model = model, messages = messages, question_context = question_context, truth_methods = truth_methods, generation_seed = generation_seed, **kwargs)
+    else:
+        return generate_with_truth_value_hf_local(model = model, messages = messages, question_context = question_context, truth_methods = truth_methods, 
+        tokenizer = tokenizer, generation_seed = generation_seed, batch_generation=batch_generation, add_generation_prompt = add_generation_prompt, continue_final_message = continue_final_message, **kwargs)
+
+
+
+            
+#TODO: remove number of generations from kwargs if exists
+def generate_with_truth_value_hf_local(model:PreTrainedModel, messages:list, question_context:str = None, truth_methods: list = [], 
+                              tokenizer: Union[PreTrainedTokenizer, PreTrainedTokenizerFast] = None, generation_seed=None, batch_generation=True, add_generation_prompt = True, continue_final_message = False,  **kwargs) -> dict:
+
+
+    tokenizer, messages = fix_tokenizer_chat(tokenizer, messages)
+    text = tokenizer.apply_chat_template(messages, tokenize = False, add_generation_prompt=add_generation_prompt, continue_final_message=continue_final_message)
+    if question_context == None:
+        question_context = ''
+        #search over last user message if exists
+        for message in messages[::-1]:
+            if message['role'] == 'user':
+                question_context = message['content']
+                break
+
+    generated_output = generate(text, model, tokenizer, **kwargs)
+    generated_text_return = generated_output['generated_text_skip_specials']
+    generated_text = generated_output['generated_text']
+    tokens = generated_output['tokens']
+    model_output = generated_output['all_ids']
+
+    #Get sampled generations to be used in truth methods
+    number_of_generations, return_text, return_logits, return_logprobs, return_attentions, return_activations = get_sampling_properties(truth_methods)
+
+    sampled_gen_dict = sample_generations_hf_local(model, text, tokenizer, generation_seed, number_of_generations=number_of_generations, 
+    return_text=return_text, return_logits=return_logits, return_logprobs=return_logprobs,return_attentions=return_attentions, return_activations=return_activations, batch_generation=batch_generation,  **kwargs)
+
+
+    # Get scores from all truth methods
+    normalized_truth_values = []
+    unnormalized_truth_values = []
+    method_spec_outputs = []
+    
+    for truth_method in truth_methods:
+        truth_values = truth_method(model=model,  input_text=text, generated_text=generated_text, question_context=question_context, all_ids=model_output, tokenizer=tokenizer, generation_seed = generation_seed, sampled_generations_dict=sampled_gen_dict, **kwargs)
+        normalized_truth_values.append(truth_values['normalized_truth_value'])
+        unnormalized_truth_values.append(truth_values['truth_value'])
+        method_spec_outputs.append(truth_values)
+
+
+    # 'all_ids': model_output.cpu(), 'generated_tokens':tokens
+    # Create TruthObject
+    truth_dict = {'generated_text':generated_text_return, 'normalized_truth_values':normalized_truth_values, 'unnormalized_truth_values':unnormalized_truth_values, 'method_specific_outputs' : method_spec_outputs, 'all_ids': model_output.cpu(), 'generated_tokens':tokens}
+
+    # Return TruthObject
+    return truth_dict
+
+
+#for api-based models, we should write a wrapper function to handle exceptions during the api call
+def generate_with_truth_value_api(model:str, messages:list, question_context:str = None, truth_methods: list = [], generation_seed=None, **kwargs) -> dict:
+    # Check if the model is an API model
+    if generation_seed is not None:
+        random.seed(generation_seed)
+
+    if type(model) == str and not model in AVAILABLE_API_MODELS:
+        raise ValueError(f"model {model} is not supported.")
+
+    requires_logprobs = False    
+    for truth_method in truth_methods:
+        if truth_method.REQUIRES_LOGPROBS:
+            requires_logprobs = True
+
+    if requires_logprobs and not model in PROB_AVAILABLE_API_MODELS:
+        raise ValueError(f"model {model} is not supported for probability requiring truth methods.")
+    
+    if question_context == None:
+        question_context = ''
+        #search over last user message if exists
+        for message in messages[::-1]:
+            if message['role'] == 'user':
+                question_context = message['content']
+                break
+
+    # Generate the main output
+    seed = kwargs.pop('seed', None)
+    if seed == None:
+        seed = random.randint(0, 1000000)
+    kwargs['seed'] = seed #a random seed is generated if seed is not specified
+
+    response = completion(
+        model=model,
+        messages=messages,
+        logprobs = requires_logprobs,
+        **kwargs
+    )
+    generated_text = response.choices[0].message['content']
+
+    logprobs = [token['logprob'] for token in response.choices[0].logprobs['content']] if requires_logprobs else None
+    generated_tokens = [token['token'] for token in response.choices[0].logprobs['content']] if requires_logprobs else None
+    
+    #Get sampled generations to be used in truth methods
+    number_of_generations, return_text, return_logits, return_logprobs, return_attentions, return_activations = get_sampling_properties(truth_methods)
+    sampled_gen_dict = sample_generations_api(model, messages, generation_seed, number_of_generations=number_of_generations, 
+    return_text=return_text, return_logits=return_logits, return_logprobs=return_logprobs,return_attentions=return_attentions, return_activations=return_activations, **kwargs)
+      
+    # Get scores from all truth methods
+    normalized_truth_values = []
+    unnormalized_truth_values = []
+    method_spec_outputs = []
+    
+    for truth_method in truth_methods:
+        truth_values = truth_method(model=model, messages=messages, generated_text=generated_text, question_context=question_context, generation_seed=generation_seed, sampled_generations_dict=sampled_gen_dict, logprobs=logprobs, generated_tokens=generated_tokens, **kwargs)
+        normalized_truth_values.append(truth_values['normalized_truth_value'])
+        unnormalized_truth_values.append(truth_values['truth_value'])
+        method_spec_outputs.append(truth_values)
+
+    # Create TruthObject
+    truth_dict = {'generated_text':generated_text, 'normalized_truth_values':normalized_truth_values, 'unnormalized_truth_values':unnormalized_truth_values, 'method_specific_outputs' : method_spec_outputs}
+
+    # Return TruthObject
+    return truth_dict
+
+
+def get_sampling_properties(truth_methods:list):
+    number_of_generations = 0
+    return_text = False
+    return_logits = False
+    return_logprobs = False
+    return_attentions = False
+    return_activations = False
+    #search over all truth methods for number of generations
+    for truth_method in truth_methods:
+        if hasattr(truth_method, 'number_of_generations') and truth_method.number_of_generations > number_of_generations:
+            number_of_generations = truth_method.number_of_generations
+        if truth_method.REQUIRES_SAMPLED_TEXT:
+            return_text = True
+        if truth_method.REQUIRES_SAMPLED_LOGITS:
+            return_logits = True
+        if truth_method.REQUIRES_SAMPLED_LOGPROBS:
+            return_logprobs = True
+        if truth_method.REQUIRES_SAMPLED_ATTENTIONS:
+            return_attentions = True
+        if truth_method.REQUIRES_SAMPLED_ACTIVATIONS:
+            return_activations = True
+    return number_of_generations, return_text, return_logits, return_logprobs, return_attentions, return_activations
+
+
+def sample_generations_hf_local(model:PreTrainedModel, input_text:str, tokenizer: Union[PreTrainedTokenizer, PreTrainedTokenizerFast], generation_seed:int=None, 
+number_of_generations:int = 0, return_text:bool = False, return_logits:bool = False, return_logprobs:bool = False, return_attentions:bool = False, return_activations:bool = False, batch_generation = False,  **kwargs):
+
+    if number_of_generations == 0 or (not return_text and not return_logprobs and not return_activations and not return_attentions and not return_logits):
+        return None
+    
+    if generation_seed is not None:
+        torch.manual_seed(generation_seed)
+        random.seed(generation_seed)
+    
+    if batch_generation == True:
+        return sample_generations_batch_hf_local(model=model, input_text=input_text, tokenizer=tokenizer, number_of_generations=number_of_generations, 
+        return_text=return_text, return_logits=return_logits, return_logprobs=return_logprobs, 
+        return_attentions=return_attentions, return_activations=return_activations, **kwargs)
+    if batch_generation == False:
+        return sample_generations_sequential_hf_local(model=model, input_text=input_text, tokenizer=tokenizer, number_of_generations=number_of_generations, 
+        return_text=return_text, return_logits=return_logits, return_logprobs=return_logprobs, 
+        return_attentions=return_attentions, return_activations=return_activations, **kwargs)
+        
+
+
+
+def sample_generations_batch_hf_local(model:PreTrainedModel, input_text:str, tokenizer: Union[PreTrainedTokenizer, PreTrainedTokenizerFast],
+number_of_generations:int = 0, return_text:bool = False, return_logits:bool = False, return_logprobs:bool = False, return_attentions:bool = False, return_activations:bool = False, return_model_output:bool = True, **kwargs):
+
+    #number_of_generations, return_text, return_logits, return_logprobs, return_attentions, return_activations = get_sampling_properties(truth_methods)
+
+    if number_of_generations == 0 or (not return_text and not return_logprobs and not return_activations and not return_attentions and not return_logits):
+        return None
+        
+    kwargs = copy.deepcopy(kwargs)
+    kwargs.pop('do_sample', None)
+    kwargs.pop('num_return_sequences', None)
+    kwargs.pop('return_dict_in_generate', None)
+    kwargs.pop('output_attentions', None)
+    kwargs.pop('output_hidden_states', None)
+    kwargs.pop('output_logits', None)
+    inputs = tokenizer(input_text, return_tensors="pt").to(model.device)
+    input_ids = inputs['input_ids']
+
+    eos_token_id = kwargs.pop("eos_token_id", None)
+
+    if eos_token_id is None:    
+        eos_token_id = model.config.eos_token_id
+
+    pad_token_id = kwargs.pop("pad_token_id", None)
+    if pad_token_id is None:
+        if type(eos_token_id) == list:
+            pad_token_id = eos_token_id[0]
+        else:
+            pad_token_id = eos_token_id
+
+    generated_texts = []
+    logits_list = []
+    logprobs = []
+    attentions_list = []
+    activations_list = []
+    tokens = []
+
+    with torch.no_grad():
+        model_output = model.generate(**inputs, num_return_sequences=number_of_generations, do_sample=True, return_dict_in_generate=True, 
+                                        output_attentions=return_attentions, output_hidden_states=return_activations, output_logits=(return_logits or return_logprobs), 
+                                        eos_token_id=eos_token_id, pad_token_id=pad_token_id, **kwargs)
+        
+        model_output.past_key_values=None
+        model_output.sequences = model_output.sequences.cpu()
+        if type(eos_token_id) == list:
+            temp = torch.stack([torch.argmax((model_output.sequences[:, len(input_ids[0]):] == eos).to(dtype=torch.int), dim=-1) for eos in eos_token_id]).T
+            indices = [torch.min(temp[i][temp[i]>0]).item() for i in range(len(temp))]
+        else:
+            indices = torch.argmax((model_output.sequences[:, len(input_ids[0]):] == eos_token_id).to(dtype=torch.int), dim=-1)
+        indices[indices==0] = model_output.sequences.shape[1] - len(input_ids[0]) -1
+        if return_text:
+            tokens = [seq[len(input_ids[0]):indices[i] + len(input_ids[0])+1].tolist() for i, seq in enumerate(model_output.sequences)]
+            generated_texts = tokenizer.batch_decode(tokens, skip_special_tokens=True)
+        if return_logprobs or return_logits:
+            logits_list = torch.stack(model_output.logits).cpu().permute(1, 0, 2)
+            model_output.logits = None
+            if return_logprobs:
+                logprobs = torch.log_softmax(logits_list, dim=-1) #logprobs for each token
+                logprobs = torch.gather(logprobs, dim=-1, index = model_output.sequences[:, len(input_ids[0]):].unsqueeze(-1))#logprobs for each token in the generated text
+                logprobs = logprobs.squeeze(-1).tolist()#convert to list
+                logprobs = [logprobs[i][:indices[i]+1] for i in range(len(logprobs))]
+            if return_logits:
+                logits_list = [logits_list[i][:indices[i]+1] for i in range(len(logits_list))]
+            else:
+                logits_list = []
+        if return_activations:
+            activations_list = [] #shape = (num gen, num token, num_layer, hidden_state_shape)
+            for i in range(number_of_generations): #generation id
+                acts = []
+                for j in range(indices[i]+1): #token id 
+                    act = []
+                    for k in range(len(model_output.hidden_states[0])): #layer id
+                        act.append(model_output.hidden_states[j][k][i].cpu())
+                    acts.append(act)
+                activations_list.append(acts)
+            model_output.hidden_states = None
+        if return_attentions:
+            attentions_list = model_output.attentions
+            for i in range(number_of_generations): #generation id
+                atts = []
+                for j in range(indices[i]+1): #token id 
+                    att = []
+                    for k in range(len(model_output.attentions[0])): #layer id
+                        att.append(model_output.attentions[j][k][i].cpu())
+                    atts.append(att)
+                attentions_list.append(atts)
+            model_output.attentions = None
+        
+        if not return_model_output:
+            model_output.sequences = None
+
+    return {"generated_texts": generated_texts, "logprobs": logprobs, "activations": activations_list, "logits":logits_list, "attentions":attentions_list, "model_outputs": model_output.sequences, "tokens":tokens}
+         
+
+
+def sample_generations_sequential_hf_local(model:PreTrainedModel, input_text:str, tokenizer: Union[PreTrainedTokenizer, PreTrainedTokenizerFast],
+ number_of_generations:int = 0, do_sample:bool=True, return_text:bool = False, return_logits:bool = False, return_logprobs:bool = False, return_attentions:bool = False, return_activations:bool = False, return_model_output:bool = True, **kwargs):
+    
+    kwargs = copy.deepcopy(kwargs)
+    kwargs.pop('do_sample', None)
+    kwargs.pop('num_return_sequences', None)
+    kwargs.pop('return_dict_in_generate', None)
+    kwargs.pop('output_attentions', None)
+    kwargs.pop('output_hidden_states', None)
+    kwargs.pop('output_logits', None)
+    inputs = tokenizer(input_text, return_tensors="pt").to(model.device)
+    input_ids = inputs['input_ids']
+
+    eos_token_id = kwargs.pop("eos_token_id", None)
+
+    if eos_token_id is None:
+        eos_token_id = model.config.eos_token_id
+
+    generated_texts = []
+    logits_list = []
+    logprobs_list = []
+    attentions_list = []
+    activations_list = []
+    model_outputs = []
+    token_lists = []
+    for i in range(number_of_generations):
+        with torch.no_grad():
+            model_output = model.generate(**inputs, num_return_sequences=1, do_sample=do_sample, return_dict_in_generate=True, 
+                                            output_attentions=return_attentions, output_hidden_states=return_activations, 
+                                            output_logits=(return_logits or return_logprobs), eos_token_id=eos_token_id, **kwargs)
+            model_output.past_key_values=None
+            model_output.sequences = model_output.sequences.cpu()
+            if return_model_output:
+                model_outputs.append(model_output.sequences)
+            if return_text:
+                tokens = model_output.sequences[0][len(input_ids[0]):]
+                generated_text = tokenizer.decode(tokens, skip_special_tokens=True)
+                generated_texts.append(generated_text)
+                token_lists.append(tokens.tolist())
+            if return_logprobs or return_logits:
+                logits = torch.cat(model_output.logits).cpu()
+                model_output.logits=None
+                if return_logprobs:
+                    logprobs = torch.log_softmax(logits, dim=-1)#logprobs for each token
+                    logprobs = torch.gather(logprobs, dim=1, index = model_output.sequences[0][len(input_ids[0]):].view(-1, 1))#logprobs for each token in the generated text
+                    logprobs = logprobs.view(-1).tolist()#convert to list
+                    logprobs_list.append(logprobs)
+                if return_logits:
+                    logits_list.append(logits)
+            if return_activations:
+                acts = []
+                for i in range(len(model_output.hidden_states)):
+                    act = []
+                    for j in range(len(model_output.hidden_states[i])):
+                        act.append(model_output.hidden_states[i][j][0].cpu())
+                    acts.append(act)
+                activations_list.append(acts)
+                model_output.hidden_states = None
+            if return_attentions:
+                atts = []
+                for i in range(len(model_output.attentions)):
+                    att = []
+                    for j in range(len(model_output.attentions[i])):
+                        att.append(model_output.attentions[i][j][0].cpu())
+                    atts.append(att)
+                attentions_list.append(atts)
+                model_output.attentions = None
+
+    return {"generated_texts": generated_texts, "logprobs": logprobs_list, "activations": activations_list, "logits":logits_list, "attentions":attentions_list, "model_outputs": model_outputs, "tokens":token_lists}
+            
+
+
+def sample_generations_api(model:str, messages:list, generation_seed:int=None, 
+number_of_generations:int = 0, return_text:bool = False, return_logits:bool = False, return_logprobs:bool = False, return_attentions:bool = False, return_activations:bool = False, **kwargs):
+    #number_of_generations, return_text, return_logits, return_logprobs, return_attentions, return_activations = get_sampling_properties(truth_methods)
+
+    if number_of_generations == 0 or (not return_text and not return_logprobs):
+        return None
+    
+
+    if generation_seed is not None:
+        random.seed(generation_seed)
+    
+    kwargs = copy.deepcopy(kwargs)
+    
+    generated_texts = []
+    logprobs_list = []
+    token_lists = []
+    for i in range(number_of_generations):
+        kwargs.pop('logprobs', None)
+        seed = kwargs.pop('seed', None) 
+        seed = random.randint(0, 1000000)
+        kwargs['seed'] = seed
+            
+        response = completion(
+            model=model,
+            messages=messages,
+            logprobs=return_logprobs,
+            **kwargs
+        )
+        
+        if return_text:
+            generated_texts.append(response.choices[0].message['content'])
+        if return_logprobs:
+            logprobs_list.append([token['logprob'] for token in response.choices[0].logprobs['content']])
+            token_lists.append([token['token'] for token in response.choices[0].logprobs['content']])
+     
+    return {"generated_texts": generated_texts, "logprobs": logprobs_list, "tokens":token_lists}

TruthTorchLM/long_form_generation/__init__.py

@@ -0,0 +1,5 @@
+from .generation import long_form_generation_with_truth_value
+from .decomposition_methods import *
+from .statement_check_methods import *
+from .evaluators import *
+from .utils import *

TruthTorchLM/long_form_generation/decomposition_methods/__init__.py

@@ -0,0 +1,8 @@
+from .decomposition_method import FactualDecompositionMethod
+from .unstructured_decomposition_api import UnstructuredDecompositionAPI
+from .unstructured_decomposition_local import UnstructuredDecompositionLocal
+from .structured_decomposition_api import StructuredDecompositionAPI
+from .structured_decomposition_local import StructuredDecompositionLocal
+
+
+__all__ = ['UnstructuredDecompositionAPI', 'UnstructuredDecompositionLocal', 'StructuredDecompositionAPI', 'StructuredDecompositionLocal']