feat: receive prompt ids

app.py

@@ -30,8 +30,7 @@ def generate_responses(model, tokenizer, prompts):
             repetition_penalty=1.0,
             max_new_tokens=512,
         )
-
-    prompt_lengths = (model_inputs.input_ids != tokenizer.pad_token_id).sum(dim=1)
+    prompt_lengths = model_inputs["attention_mask"].sum(dim=1)
     generated_ids = [
         output_ids[length:] for length, output_ids in zip(prompt_lengths, generated_ids)
     ]
@@ -55,13 +54,16 @@ def classify_pairs(model, tokenizer, prompts, responses):
 
 
 @spaces.GPU()
-def generate(prompts: list[str]) -> list[dict[str, str | float]]:
+def generate(submission: list[dict[str, str]]) -> list[dict[str, str | float]]:
+    ids = [s["id"] for s in submission]
+    prompts = [s["prompt"] for s in submission]
+
     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)
+        {"id": id, "prompt": prompt, "response": response, "score": score}
+        for id, prompt, response, score in zip(ids, prompts, responses, scores)
     ]