Shuggu commited on
Commit
c42ab87
1 Parent(s): 5611a81
Files changed (1) hide show
  1. app.py +11 -0
app.py ADDED
@@ -0,0 +1,11 @@
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from transformers import AutoTokenizer
2
+ from safe_rlhf.models import AutoModelForScore
3
+
4
+ model = AutoModelForScore.from_pretrained('PKU-Alignment/beaver-7b-v1.0-reward', device_map='auto')
5
+ tokenizer = AutoTokenizer.from_pretrained('PKU-Alignment/beaver-7b-v1.0-reward', use_fast=False)
6
+
7
+ input = 'BEGINNING OF CONVERSATION: USER: hello ASSISTANT:Hello! How can I help you today?'
8
+
9
+ input_ids = tokenizer(input, return_tensors='pt')
10
+ output = model(**input_ids)
11
+ print(output)