Views
No views yet
Summarise the article:\\n{article} |||\\n{positive_sample}Summarise the article:\\n{article} |||\\n1from peft import PeftModel, PeftConfig
2from transformers import AutoModelForCausalLM, AutoTokenizer
3import torch
4
5source_model_id = "NbAiLab/nb-gpt-j-6B"
6peft_model_id = "NorGLM/NbAiLab-6B-summarization-peft"
7
8config = PeftConfig.from_pretrained(peft_model_id)
9model = AutoModelForCausalLM.from_pretrained(source_model_id, device_map='balanced')
10
11tokenizer_max_len = 2048
12tokenizer_config = {'pretrained_model_name_or_path': source_model_id,
13 'max_len': tokenizer_max_len}
14tokenizer = tokenizer = AutoTokenizer.from_pretrained(**tokenizer_config)
15tokenizer.pad_token = tokenizer.eos_token
16
17model = PeftModel.from_pretrained(model, peft_model_id)1def generate_texts(model, tokenizer, prompts, max_seq_length=200, do_sample=True, top_p=0.95, top_k=10):
2 # prompts are a list of news articles
3 results = []
4 cnt = 0
5 for prompt in prompts:
6 cnt += 1
7 pro_len = len(prompt.split())
8 if pro_len>1024:
9 results.append('')
10 continue
11
12 prompt = 'Summarise the article:\\n' + prompt + ' |||\\n'
13
14 model_inputs = tokenizer(prompt, return_tensors='pt').to(torch_device)
15 output = model.generate(**model_inputs, do_sample=False, max_new_tokens=max_seq_length)
16 result = tokenizer.decode(output[0], skip_special_tokens=True)
17 result = result.split("|||\\n")[-1]
18 results.append(result)
19 return results
20
21print("--LOADING EVAL DATAS---")
22eval_data = load_dataset("NorGLM/NO-CNN-DailyMail", data_files="test.csv")
23prompts = eval_data['train']['article']
24positive_samples = eval_data['train']['positive_sample']
25
26print("--MAKING PREDICTIONS---")
27model.eval()
28
29output_file = <output file name>
30with torch.no_grad():
31 results = generate_texts(model, tokenizer, prompts)
32
33df = pd.DataFrame({'article':prompts, 'generated_text':results, 'positive_sample':positive_samples})
34
35print("Save results to csv file...")
36df.to_csv(output_file)
37