Update app.py

app.py

@@ -1,6 +1,6 @@
 import spaces
 import gradio as gr
-from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
+from transformers import AutoModelForCausalLM, AutoTokenizer, AutoModelForSequenceClassification, BitsAndBytesConfig
 import torch
 from gradio_rich_textbox import RichTextbox
 
@@ -16,15 +16,35 @@ model_path = "NousResearch/Genstruct-7B"
 tokenizer = AutoTokenizer.from_pretrained(model_path)
 quantization_config = BitsAndBytesConfig(load_in_8bit=True)
 model = AutoModelForCausalLM.from_pretrained(model_path, device_map='cuda', quantization_config=quantization_config)
+rm_tokenizer = AutoTokenizer.from_pretrained('OpenAssistant/reward-model-deberta-v3-large-v2')
+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):
-    input_text = {'title': usertitle, 'content': content} 
+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_text = tokenizer.decode(model.generate(inputs, max_new_tokens=max_length, temperature=temperature, do_sample=True)[0]).strip().split(tokenizer.eos_token)[0]
-    # split_text = generated_text.split(tokenizer.eos_token)[0]
+    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)
 
-    return generated_text
+    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)
+        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
+
+    best_sequence = max(decoded_sequences, key=score)
+
+    return best_sequence
 
 def gradio_app():
     with gr.Blocks() as demo: