Views
No views yet
distilroberta-base model (https://huggingface.co/distilroberta-base) as a base.1from transformers import (
2 AutoModelForSequenceClassification,
3 AutoTokenizer,
4 pipeline
5)
6
7reward_model = AutoModelForSequenceClassification.from_pretrained(
8 cambioml/rlhf_reward_model,
9 num_labels=1,
10 # torch_dtype=torch.bfloat16,
11 load_in_8bit=True,
12 device_map={"": Accelerator().process_index}
13)
14
15reward_tokenizer = AutoTokenizer.from_pretrained(cambioml/rlhf_reward_model)
16reward_tokenizer.pad_token = reward_tokenizer.eos_token
17
18
19reward_kwargs = {
20 "return_all_scores": True,
21 "function_to_apply": "none",
22 "batch_size": 32,
23 "truncation": True,
24 "max_length": 138
25}
26
27reward_pipe = pipeline(
28 "sentiment-analysis",
29 model=reward_model,
30 model_kwargs=reward_kwargs,
31 tokenizer=reward_tokenizer,
32 return_token_type_ids=False,
33)