Views
No views yet
1from transformers import AutoModelForCausalLM, AutoConfig, AutoTokenizer
2import torch
3
4device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
5model_name = "Aconoya/Nono_instruct_neo-125m_dpo"
6model = AutoModelForCausalLM.from_pretrained(model_name)
7model.gradient_checkpointing_enable()
8model = model.to(device.type)
9tokenizer = AutoTokenizer.from_pretrained(model_name)
10tokenizer.pad_token = tokenizer.eos_token
11
12ender_string = '<endofturn>'
13system_string='<system>'
14user_string='<user>'
15assistant_string='<assistant>'
16
17prompt = ['Hello! How are you?', '¡Hola!, ¿Cómo estás?', '¿Qué es un perro?', 'What is a dog?']
18prompt = choice(prompt)
19formatted_prompt = system_string + 'You are a digital assistant.' + ender_string + '\n' + user_string + prompt + ender_string + '\n' + assistant_string
20model_input = tokenizer.encode(formatted_prompt, return_tensors='pt').to(device)
21generated_ids = model.generate(input_ids=model_input, pad_token_id=tokenizer.eos_token_id, max_new_tokens=50)
22generated_text = tokenizer.decode(generated_ids[:, model_input.shape[-1]:][0], skip_special_tokens=True)
23print('Prompt:', prompt)
24print("Response: '{}'".format(generated_text))