Update app.py

app.py

@@ -20,27 +20,18 @@ rm_tokenizer = AutoTokenizer.from_pretrained('OpenAssistant/reward-model-deberta
 rm_model = AutoModelForSequenceClassification.from_pretrained('OpenAssistant/reward-model-deberta-v3-large-v2', torch_dtype=torch.bfloat16)
 
 @spaces.GPU
-def generate_text(usertitle, content, max_length, temperature, N=3):
-    input_text = {'title': usertitle, 'content': content}
-    inputs = tokenizer.apply_chat_template(input_text, return_tensors='pt').cuda()
-    generated_sequences = model.generate(inputs, max_new_tokens=max_length, temperature=temperature, num_return_sequences=N, do_sample=True)
-    decoded_sequences = tokenizer.batch_decode(generated_sequences, skip_special_tokens=True)
-
-    def extract_pair(resp):
-        try:
-            response = resp.split('[[[Content]]]')[1]
-            inst, resp = resp.split('[[[User]]]')[:2]
-        return inst.strip(), resp.strip()
-        except ValueError:
-            return "", ""
-
-    def score(resp):
-        inst, resp = extract_pair(resp)
+def generate_text(usertitle, content, max_length, temperature, N=5):
+    input_text = f"title: {usertitle}\ncontent: {content}"
+    inputs = tokenizer(input_text, return_tensors='pt').to('cuda')
+    generated_sequences = model.generate(inputs['input_ids'], max_length=max_length, temperature=temperature, num_return_sequences=N, do_sample=True)
+    decoded_sequences = [tokenizer.decode(g, skip_special_tokens=True) for g in generated_sequences]
+
+    def score(sequence):
+        inputs = rm_tokenizer(sequence, return_tensors='pt', padding=True, truncation=True, max_length=max_length).to('cuda')
         with torch.no_grad():
-            inputs = rm_tokenizer(inst, resp, return_tensors='pt', truncation=True, max_length=512)
-            logits = rm_model(**inputs).logits
-            score = float(logits[0].cpu())
-            return score
+            outputs = rm_model(**inputs)
+            logits = outputs.logits
+            return logits[0][1].item()
 
     best_sequence = max(decoded_sequences, key=score)