feat: update output format

app.py

@@ -1,45 +1,72 @@
 import gradio as gr
 import spaces
-from transformers import pipeline
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import (
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    AutoModelForSequenceClassification,
+)
 import torch
-import json
 
-model_name = "sapienzanlp/Minerva-7B-instruct-v1.0"
-model = AutoModelForCausalLM.from_pretrained(
-    model_name,
-    torch_dtype=torch.bfloat16,
-    device_map="auto"
+chat_model_name = "sapienzanlp/Minerva-7B-instruct-v1.0"
+chat_model = AutoModelForCausalLM.from_pretrained(
+    chat_model_name, torch_dtype=torch.bfloat16, device_map="auto"
+)
+chat_tokenizer = AutoTokenizer.from_pretrained(chat_model_name)
+
+moderator_model_name = "saiteki-kai/QA-DeBERTa-v3-large"
+moderator_model = AutoModelForSequenceClassification.from_pretrained(
+    moderator_model_name, device_map="auto"
 )
-tokenizer = AutoTokenizer.from_pretrained(model_name)
+moderator_tokenizer = AutoTokenizer.from_pretrained(moderator_model_name)
 
-classifier = pipeline("text-classification", model="saiteki-kai/QA-DeBERTa-v3-large")
 
-@spaces.GPU()
-def generate(prompts: list[str]) -> list[tuple[str, dict[str, float]]]:
+def generate_responses(model, tokenizer, prompts):
     messages = [[{"role": "user", "content": message}] for message in prompts]
 
     texts = tokenizer.apply_chat_template(
-        messages,
-        tokenize=False,
-        add_generation_prompt=True
-    )
-    model_inputs = tokenizer(texts, padding=True, return_tensors="pt").to(model.device)
-    generated_ids = model.generate(
-        **model_inputs,
-        do_sample=False,
-        temperature=0,
-        repetition_penalty=1.0,
-        max_new_tokens=512,
+        messages, tokenize=False, add_generation_prompt=True
     )
 
+    with torch.inference_mode():
+        model_inputs = tokenizer(texts, padding=True, return_tensors="pt").to(model.device)
+        generated_ids = model.generate(
+            **model_inputs,
+            do_sample=False,
+            temperature=0,
+            repetition_penalty=1.0,
+            max_new_tokens=512,
+        )
+
     prompt_lengths = (model_inputs.input_ids != tokenizer.pad_token_id).sum(dim=1)
     generated_ids = [
         output_ids[length:] for length, output_ids in zip(prompt_lengths, generated_ids)
     ]
     responses = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)
 
-    return list(zip(responses, classifier([text + "[SEP]" + response for text, response in zip(texts, responses)])))
+    return responses
+
+
+def classify_pairs(model, tokenizer, prompts, responses):
+    texts = [
+        prompt + "[SEP]" + response for prompt, response in zip(prompts, responses)
+    ]
+
+    with torch.inference_mode():
+        input_ids = tokenizer(texts, padding=True, max_length=512).to(model.device)
+        outputs = model(**input_ids)
+
+    return outputs
+
+
+@spaces.GPU()
+def generate(prompts: list[str]) -> list[dict[str, str | float]]:
+    responses = generate_responses(chat_model, chat_tokenizer, prompts)
+    scores = classify_pairs(moderator_model, moderator_tokenizer, prompts, responses)
+
+    return [
+        {"prompt": prompt, "response": response, "score": score}
+        for prompt, response, score in zip(prompts, responses, scores)
+    ]
 
 
 with gr.Blocks() as demo:
@@ -47,4 +74,4 @@ with gr.Blocks() as demo:
     gr.api(generate, api_name="scores", batch=False)
 
 demo.queue()
-demo.launch()
+demo.launch()