dna_gpt.py

# Copyright (c) Guangsheng Bao.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
import os.path

import numpy as np
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
import re
import torch
import tqdm
import argparse
import json
from data_builder import load_data, save_data
from metrics import get_roc_metrics, get_precision_recall_metrics
from model import load_tokenizer, load_model, get_model_fullname, from_pretrained
from data_builder import load_data
from model import load_tokenizer, load_model
from metrics import get_roc_metrics, get_precision_recall_metrics
import custom_datasets

class PrefixSampler:
    def __init__(self, args):
        self.args = args
        self.base_tokenizer = load_tokenizer(args.base_model_name, args.dataset, args.cache_dir)
        self.base_model = load_model(args.base_model_name, args.device, args.cache_dir)

    def _sample_from_model(self, texts, min_words=55, truncate_ratio=0.5):
        # encode each text as a list of token ids
        if self.args.dataset == 'pubmed':
            pubmed_sep = ' Answer:'
            texts = [t[:t.index(pubmed_sep) + len(pubmed_sep)] for t in texts]
            all_encoded = self.base_tokenizer(texts, return_tensors="pt", padding=True).to(self.args.device)
        else:
            texts = [t.split(' ') for t in texts]
            texts = [' '.join(t[: int(len(t) * truncate_ratio)]) for t in texts]
            all_encoded = self.base_tokenizer(texts, return_tensors="pt", padding=True).to(self.args.device)

        self.base_model.eval()
        decoded = ['' for _ in range(len(texts))]

        # sample from the model until we get a sample with at least min_words words for each example
        # this is an inefficient way to do this (since we regenerate for all inputs if just one is too short), but it works
        tries = 0
        m = 0
        while m < min_words:
            if tries != 0:
                print()
                print(f"min words: {m}, needed {min_words}, regenerating (try {tries})")

            sampling_kwargs = {'temperature': self.args.temperature}
            if self.args.do_top_p:
                sampling_kwargs['top_p'] = self.args.top_p
            elif self.args.do_top_k:
                sampling_kwargs['top_k'] = self.args.top_k
            min_length = 50 if self.args.dataset in ['pubmed'] else 150
            outputs = self.base_model.generate(**all_encoded, min_length=min_length, max_length=200, do_sample=True,
                                               **sampling_kwargs, pad_token_id=self.base_tokenizer.eos_token_id,
                                               eos_token_id=self.base_tokenizer.eos_token_id)
            decoded = self.base_tokenizer.batch_decode(outputs, skip_special_tokens=True)
            m = min(len(x.split()) for x in decoded)
            tries += 1

        return decoded

    def generate_samples(self, raw_data, batch_size):
        # trim to shorter length
        def _trim_to_shorter_length(texta, textb):
            # truncate to shorter of o and s
            shorter_length = min(len(texta.split(' ')), len(textb.split(' ')))
            texta = ' '.join(texta.split(' ')[:shorter_length])
            textb = ' '.join(textb.split(' ')[:shorter_length])
            return texta, textb

        def _truncate_to_substring(text, substring, idx_occurrence):
            # truncate everything after the idx_occurrence occurrence of substring
            assert idx_occurrence > 0, 'idx_occurrence must be > 0'
            idx = -1
            for _ in range(idx_occurrence):
                idx = text.find(substring, idx + 1)
                if idx == -1:
                    return text
            return text[:idx]

        data = {
            "original": [],
            "sampled": [],
        }

        assert len(raw_data) % batch_size == 0
        for batch in range(len(raw_data) // batch_size):
            print('Generating samples for batch', batch, 'of', len(raw_data) // batch_size)
            original_text = raw_data[batch * batch_size:(batch + 1) * batch_size]
            sampled_text = self._sample_from_model(original_text, min_words=30 if self.args.dataset in ['pubmed'] else 55, truncate_ratio=self.args.truncate_ratio)

            for o, s in zip(original_text, sampled_text):
                if self.args.dataset == 'pubmed':
                    s = _truncate_to_substring(s, 'Question:', 2)
                    o = o.replace(custom_datasets.SEPARATOR, ' ')

                o, s = _trim_to_shorter_length(o, s)

                # add to the data
                data["original"].append(o)
                data["sampled"].append(s)

        return data

def get_likelihood(logits, labels, pad_index):
    labels = labels.unsqueeze(-1) if labels.ndim == logits.ndim - 1 else labels
    lprobs = torch.log_softmax(logits, dim=-1)
    log_likelihood = lprobs.gather(dim=-1, index=labels)
    mask = labels != pad_index
    log_likelihood = (log_likelihood * mask).sum(dim=1) / mask.sum(dim=1)
    return log_likelihood.squeeze(-1)

def get_log_prob(sampler, text):
    tokenized = sampler.base_tokenizer(text, return_tensors="pt", padding=True).to(sampler.args.device)
    labels = tokenized.input_ids[:, 1:]
    with torch.no_grad():
        logits_score = sampler.base_model(**tokenized).logits[:, :-1]
        return get_likelihood(logits_score, labels, sampler.base_tokenizer.pad_token_id)

def get_log_probs(sampler, texts):
    batch_size = sampler.args.batch_size
    batch_lprobs = []
    for batch in range(len(texts) // batch_size):
        tokenized = sampler.base_tokenizer(texts[batch * batch_size:(batch + 1) * batch_size], return_tensors="pt", padding=True).to(sampler.args.device)
        labels = tokenized.input_ids[:, 1:]
        with torch.no_grad():
            logits_score = sampler.base_model(**tokenized).logits[:, :-1]
            lprobs = get_likelihood(logits_score, labels, sampler.base_tokenizer.pad_token_id)
            batch_lprobs.append(lprobs)
    return torch.cat(batch_lprobs, dim=0)

def get_regen_samples(sampler, text):
    data = [text] * sampler.args.regen_number
    data = sampler.generate_samples(data, batch_size=sampler.args.batch_size)
    return data['sampled']

def get_dna_gpt(sampler, text):
    lprob = get_log_prob(sampler, text)
    regens = get_regen_samples(sampler, text)
    lprob_regens = get_log_probs(sampler, regens)
    wscore = lprob[0] - lprob_regens.mean()
    return wscore.item()

def experiment(args):
    sampler = PrefixSampler(args)
    # load data
    data = load_data(args.dataset_file)
    n_samples = len(data["sampled"])
    # evaluate criterion
    name = "dna_gpt"
    criterion_fn = get_dna_gpt

    torch.manual_seed(args.seed)
    np.random.seed(args.seed)
    results = []
    for idx in tqdm.tqdm(range(n_samples), desc=f"Computing {name} criterion"):
        original_text = data["original"][idx]
        sampled_text = data["sampled"][idx]
        # original text
        original_crit = criterion_fn(sampler, original_text)
        # sampled text
        sampled_crit = criterion_fn(sampler, sampled_text)
        # result
        results.append({"original": original_text,
                        "original_crit": original_crit,
                        "sampled": sampled_text,
                        "sampled_crit": sampled_crit})

    # compute prediction scores for real/sampled passages
    predictions = {'real': [x["original_crit"] for x in results],
                   'samples': [x["sampled_crit"] for x in results]}
    fpr, tpr, roc_auc = get_roc_metrics(predictions['real'], predictions['samples'])
    p, r, pr_auc = get_precision_recall_metrics(predictions['real'], predictions['samples'])
    print(f"Criterion {name}_threshold ROC AUC: {roc_auc:.4f}, PR AUC: {pr_auc:.4f}")
    # results
    results_file = f'{args.output_file}.{name}.json'
    results = { 'name': f'{name}_threshold',
                'info': {'n_samples': n_samples},
                'predictions': predictions,
                'raw_results': results,
                'metrics': {'roc_auc': roc_auc, 'fpr': fpr, 'tpr': tpr},
                'pr_metrics': {'pr_auc': pr_auc, 'precision': p, 'recall': r},
                'loss': 1 - pr_auc}
    with open(results_file, 'w') as fout:
        json.dump(results, fout)
        print(f'Results written into {results_file}')

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--output_file', type=str, default="./exp_test/results/pubmed_davinci")
    parser.add_argument('--dataset', type=str, default="pubmed")
    parser.add_argument('--dataset_file', type=str, default="./exp_test/data/pubmed_davinci")
    parser.add_argument('--truncate_ratio', type=float, default=0.5)
    parser.add_argument('--regen_number', type=int, default=10)
    parser.add_argument('--base_model_name', type=str, default="gpt2")
    parser.add_argument('--batch_size', type=int, default=10)
    parser.add_argument('--do_top_k', action='store_true')
    parser.add_argument('--top_k', type=int, default=40)
    parser.add_argument('--do_top_p', action='store_true')
    parser.add_argument('--top_p', type=float, default=0.96)
    parser.add_argument('--temperature', type=float, default=1.0)
    parser.add_argument('--seed', type=int, default=0)
    parser.add_argument('--device', type=str, default="cuda")
    parser.add_argument('--cache_dir', type=str, default="../cache")
    args = parser.parse_args()

    experiment(args)