rand.py

# 不共享ice，随机样本作为ice

from sentence_transformers import SentenceTransformer
from transformers import AutoTokenizer, GPT2Tokenizer, AutoModelForCausalLM, GPT2Model
import faiss
from typing import Dict
from transformers import pipeline
from torch.utils.data import DataLoader
from tqdm.auto import tqdm
from torch.utils.data import Dataset
import evaluate
import torch
from datasets import load_dataset
import yaml
import numpy as np
from sklearn.metrics.pairwise import cosine_similarity

class ListDataset(Dataset):
     def __init__(self, original_list):
        self.original_list = original_list
        
     def __len__(self):
        return len(self.original_list)

     def __getitem__(self, i):
        return self.original_list[i]
    
    
def get_embeddings(corpus, model_name="all-mpnet-base-v2"):
    model = SentenceTransformer(model_name)
    embeddings = model.encode(corpus, show_progress_bar=True, convert_to_numpy=True)
    return embeddings

def create_index(embeddings):
    index = faiss.IndexFlatIP(768)
    index.add(embeddings)
    return index

def knn_search(index, corpus, ice_num=8):
    print("Embedding search data...")
    query_embeddings = get_embeddings(corpus)
    
    _, neighbours = index.search(query_embeddings, ice_num)
    return neighbours    


def generate_item(template, dataset, ice_id, label_map: Dict, input_column="text", label_column="label"):
    label = dataset[label_column][ice_id]
    prompt = template.format(text=dataset[input_column][ice_id], verb=label_map[label])
    return prompt
    

def generate_prompts(template, dataset: Dict, ice_ids, label_map: Dict, split="test", input_column="text", label_column="label"):
    prompts = []
    
    # construct a prompt for every test data point
    for idx, neighbours in enumerate(tqdm(ice_ids)):
        prompt = "\n".join([generate_item(template, dataset["train"], i, label_map, input_column, label_column) for i in neighbours]) + "\n"
        prompt += template.format(text=dataset[split][input_column][idx], verb="")
        prompts.append(prompt)
    
    return prompts

def data(prompts):
    for prompt in prompts:
        yield prompt

def evaluate_result(outputs, labels):
    metric = evaluate.load("accuracy")
    results = metric.compute(references=labels, predictions=outputs)
    return results

def inference1(model_name, prompts):
    pipe = pipeline("text-generation", model=model_name, device="cuda", batch_size=8, return_full_text=False, max_length=500)
    model = AutoModelForCausalLM.from_pretrained(model_name)
    pipe.tokenizer.pad_token_id = model.config.eos_token_id
    pipe.tokenizer.padding_side = "left"
    
    dataset = ListDataset(prompts)
    
    outputs = []
    
    for out in tqdm(pipe(dataset, pad_token_id=pipe.tokenizer.eos_token_id)):
        s = out[0]["generated_text"].split("\n")[0]
        if s not in ["negative", "positive"]:
            
            import pdb; pdb.set_trace()
        outputs.append(s)
    return outputs

# def rerank(ids):
    
#     # this reranks the example order
#     ice_num = ids.shape[-1]
    
#     i, j = 0, ice_num - 1
#     counter = 0
#     permutation = np.zeros(ice_num).astype(int)
    
#     while i < j:
#         permutation[i] = counter
#         counter += 1
#         permutation[j] = counter
#         counter += 1
#         i += 1
#         j -= 1
#     if i == j:
#         permutation[i] = counter
#     permutation = np.arange(ice_num)[::-1]
#     print("rerank permutation: ", permutation)
#     ids = ids[:, permutation]
#     return ids

def rerank(ids, similarity):
    order = np.zeros_like(ids)
    
    for i in range(ids.shape[0]):
        neighbour = ids[i]
        sim = similarity[i][neighbour]
        index = np.argsort(sim) # 升序排列
        index = index[::-1]
        order[i] = neighbour[index]
        # order[i] = index
    return order

def inference(model_name, prompts, batch_size=8):
    device = "cuda"
    
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    tokenizer.pad_token = tokenizer.eos_token
    tokenizer.pad_token_id = tokenizer.eos_token_id
    tokenizer.padding_side = "left"
    
    model = AutoModelForCausalLM.from_pretrained(model_name).to(device)
    
    labels = []
    print("len prompts: ", len(prompts))
    num_batches = len(prompts) // batch_size if len(prompts) % batch_size == 0 else len(prompts) // batch_size + 1
    
    for i in tqdm(range(num_batches)):
        batch = prompts[i*batch_size: (i+1)*batch_size]
        
        tokens = tokenizer(batch, return_tensors='pt', padding=True)
        input_ids = tokens["input_ids"].to(device)
        attention_mask = tokens["attention_mask"].to(device)
        output = model.generate(input_ids=input_ids, attention_mask=attention_mask, max_new_tokens=1, pad_token_id=tokenizer.eos_token_id,
                                       eos_token_id=tokenizer.eos_token_id)
        length = input_ids.shape[1]
        output = output.detach().cpu().numpy()
        
        pred = list(output[:, length])
        label = [tokenizer.decode(p, skip_special_tokens=True) for p in pred]
        labels.extend(label)
    torch.cuda.empty_cache()
    return labels
        
    # for prompt in prompts:
    #     input_ids = tokenizer.encode(prompt, return_tensors='pt').to(device)
    #     output = model.generate(input_ids, max_new_tokens=1)
    #     generated_text = tokenizer.decode(output[0], skip_special_tokens=True)
    #     prediction = generated_text.removeprefix(prompt)
    #     # print("Prompt: ", prompt)
    #     print("Prediction text: ", prediction)
    #     if prediction not in ["positive", "negative"]:
    #         import pdb; pdb.set_trace()

    
    # tokenize prompts
    # input_ids = tokenizer(prompts, return_tensors="pt", padding=True, max_length=None).to(device)
    
    # output = model(**input_ids)

def random_neighbours(train, test, ice_num=8):
    neighbours = np.zeros((test, ice_num), dtype=int)
    
    fudge = np.random.choice(train, ice_num, replace=False)
    for i in range(test):
        neighbours[i] = np.random.choice(train, ice_num, replace=False)
    return neighbours

def perm(ids):
    # this reranks the example order
    ice_num = ids.shape[-1]
    
    i, j = 0, ice_num - 1
    counter = 0
    permutation = np.zeros(ice_num).astype(int)
    
    while i < j:
        permutation[i] = counter
        counter += 1
        permutation[j] = counter
        counter += 1
        i += 1
        j -= 1
    if i == j:
        permutation[i] = counter
    # permutation = np.arange(ice_num)[::-1]
    print("rerank permutation: ", permutation)
    ids = ids[:, permutation]
    return ids


def perm1(ids):
    # this reranks the example order
    ice_num = ids.shape[-1]
    
    i, j = 0, ice_num - 1
    counter = ice_num - 1
    permutation = np.zeros(ice_num).astype(int)
    
    while i < j:
        permutation[i] = counter
        counter -= 1
        permutation[j] = counter
        counter -= 1
        i += 1
        j -= 1
    if i == j:
        permutation[i] = counter
    # permutation = np.arange(ice_num)[::-1]
    print("rerank permutation: ", permutation)
    ids = ids[:, permutation]
    return ids

def reverse_dict(d):
    return {v: k for k, v in d.items()}

if __name__ == "__main__":
    with open("config.yaml", "r") as f:
        config = yaml.load(f, Loader=yaml.FullLoader)
        print(config)

    # Loading dataset from huggingface
    dataset = load_dataset(config["dataset"])
    train_corpus = dataset["train"]["text"]
    test_corpus = dataset[config["split"]]["text"]
    train_embeddings = get_embeddings(train_corpus)
    test_embeddings = get_embeddings(test_corpus)
    sim = cosine_similarity(train_embeddings, test_embeddings).T
    print("similarity matrix shape: ", sim.shape)
    
    neighbours = random_neighbours(len(train_corpus), len(test_corpus), config["ice_num"])
    # neighbours = rerank(neighbours, sim)
    # neighbours = perm(neighbours)    
    
    template = "Movie Review:{text}\nSentiment:{verb}"
    label_map = config["label_map"]
    map_label = reverse_dict(label_map)
    prompts = generate_prompts(template, dataset, neighbours, label_map, config["split"], config["input_column"], config["output_column"])
    print(prompts[:5])
    
    # import pickle
    # with open("icl_out.bin", "wb") as f:
    #     pickle.dump(neighbours, f)
    predictions = inference(config["model"], prompts, batch_size=config["batch_size"])  
    
    pred = list(map(lambda x: map_label[x], predictions))
    test_labels = dataset[config["split"]][config["output_column"]]
    print(evaluate_result(pred, test_labels))
    
    # print("相似度按我们的来")
    # neighbours1 = perm(neighbours)
    # prompts = generate_prompts(template, dataset, neighbours1, label_map, config["split"], config["input_column"], config["output_column"])
    # predictions = inference(config["model"], prompts)  
    
    # pred = list(map(lambda x: map_label[x], predictions))
    # test_labels = dataset[config["split"]][config["output_column"]]
    # print(evaluate_result(pred, test_labels))
    
    # print("相似度按我们的反着来")
    # neighbours2 = perm1(neighbours)
    # prompts = generate_prompts(template, dataset, neighbours2, label_map, config["split"], config["input_column"], config["output_column"])
    # predictions = inference(config["model"], prompts)  
    
    # pred = list(map(lambda x: map_label[x], predictions))
    # test_labels = dataset[config["split"]][config["output_column"]]
    # print(evaluate_result(pred, test_labels))