NbAiLab-6B-summarization-peft

Model Card

NbAiLab-6B-summarization-peft is trained on top of NbAiLab/nb-gpt-j-6B model on NO-CNN-DailyMail dataset.

Prompt format:

Summarise the article:\\n{article} |||\\n{positive_sample}

Inference prompt:

Summarise the article:\\n{article} |||\\n

Run the Model

from peft import PeftModel, PeftConfig
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

source_model_id = "NbAiLab/nb-gpt-j-6B"
peft_model_id = "NorGLM/NbAiLab-6B-summarization-peft"

config = PeftConfig.from_pretrained(peft_model_id)
model = AutoModelForCausalLM.from_pretrained(source_model_id, device_map='balanced')

tokenizer_max_len = 2048
tokenizer_config = {'pretrained_model_name_or_path': source_model_id,
                            'max_len': tokenizer_max_len}
tokenizer = tokenizer = AutoTokenizer.from_pretrained(**tokenizer_config)
tokenizer.pad_token = tokenizer.eos_token

model = PeftModel.from_pretrained(model, peft_model_id)

Inference on test set

Load the model to evaluate on the test set of NO-CNN-DailyMail dataset:

def generate_texts(model, tokenizer, prompts, max_seq_length=200, do_sample=True, top_p=0.95, top_k=10):
    # prompts are a list of news articles
    results = []
    cnt = 0
    for prompt in prompts:
        cnt += 1
        pro_len = len(prompt.split())
        if pro_len>1024:
            results.append('')
            continue

        prompt = 'Summarise the article:\\n' + prompt + ' |||\\n'

        model_inputs = tokenizer(prompt, return_tensors='pt').to(torch_device)
        output = model.generate(**model_inputs, do_sample=False, max_new_tokens=max_seq_length)
        result = tokenizer.decode(output[0], skip_special_tokens=True)
        result = result.split("|||\\n")[-1]
        results.append(result)
    return results

print("--LOADING EVAL DATAS---")
eval_data = load_dataset("NorGLM/NO-CNN-DailyMail", data_files="test.csv")
prompts = eval_data['train']['article']
positive_samples = eval_data['train']['positive_sample']

print("--MAKING PREDICTIONS---")
model.eval()

output_file = <output file name>
with torch.no_grad():
    results = generate_texts(model, tokenizer, prompts)

df = pd.DataFrame({'article':prompts, 'generated_text':results, 'positive_sample':positive_samples})

print("Save results to csv file...")
df.to_csv(output_file)

Note

More training details will be released soon!