Views
No views yet
1DEVICE = torch.device('cuda')
2
3tokenizer = AutoTokenizer.from_pretrained(llm-course-hw2-reward-model)
4reward_model = AutoModelForSequenceClassification.from_pretrained(llm-course-hw2-reward-model, num_labels = 1)
5reward_model.config.pad_token_id = tokenizer.pad_token_id
6reward_model = reward_model.to(DEVICE)
7reward_model.eval()
8
9inputs_chosen = tokenizer.apply_chat_template('Any text'], tokenize=False)
10inputs_chosen = tokenizer(inputs_chosen, return_tensors="pt").to(DEVICE)
11
12
13score_chosen = reward_model(**inputs_chosen).logits[0].cpu().detach()
14print(score_chosen)
15
16#0.2231@misc{vonwerra2022trl,
2 title = {{TRL: Transformer Reinforcement Learning}},
3 author = {Leandro von Werra and Younes Belkada and Lewis Tunstall and Edward Beeching and Tristan Thrush and Nathan Lambert and Shengyi Huang and Kashif Rasul and Quentin Gallouédec},
4 year = 2020,
5 journal = {GitHub repository},
6 publisher = {GitHub},
7 howpublished = {\url{https://github.com/huggingface/trl}}
8}