Update app.py

app.py

@@ -51,42 +51,21 @@ print("\n=== Model Loading ===")
 chat_model_name = "sapienzanlp/Minerva-7B-instruct-v1.0"
 cls_model_name = "saiteki-kai/QA-DeBERTa-v3-large-binary-3"
 
-chat_model = AutoModelForCausalLM.from_pretrained(chat_model_name, dtype=torch.bfloat16)
-cls_model = AutoModelForSequenceClassification.from_pretrained(cls_model_name, dtype=torch.bfloat16)
+pipeline = transformers.pipeline(
+    model=chat_model_name,
+    model_kwargs={"dtype": torch.bfloat16},
+    device=device,
+)
 
-chat_tokenizer = AutoTokenizer.from_pretrained(chat_model_name)
+cls_model = AutoModelForSequenceClassification.from_pretrained(cls_model_name, dtype=torch.bfloat16)
 cls_tokenizer = AutoTokenizer.from_pretrained(cls_model_name)
-
-chat_model = chat_model.to(device)  # type: ignore
 cls_model = cls_model.to(device)
 
-def generate_responses(model, tokenizer, prompts):
-    texts = []
-    for prompt in prompts:
-        message = [{"role": "user", "content": prompt}]
-        text = tokenizer.apply_chat_template(message, tokenize=False, add_generation_prompt=True)
-        texts.append(text)
-
-    print(texts[0])
-
-    model_inputs = tokenizer(texts, padding=True, truncation=True, max_length=512, return_tensors="pt").to(model.device)
-    
-    print(tokenizer.decode(model_inputs["input_ids"][0]))
-    
-    with torch.inference_mode():
-        generated_ids = model.generate(
-            **model_inputs,
-            do_sample=False,
-            temperature=0,
-            repetition_penalty=1.1,
-            max_new_tokens=512,
-        )
-    prompt_lengths = model_inputs["attention_mask"].sum(dim=1)
-    generated_ids = [output_ids[length:] for length, output_ids in zip(prompt_lengths, generated_ids)]
-    print(tokenizer.decode(generated_ids[0], skip_special_tokens=False))
-    responses = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)
+def generate_responses(pipeline, tokenizer, prompts):    
+    messages = [[{"role": "user", "content": prompt}] for prompt in prompts]
+    responses = pipeline(messages, do_sample=False, max_new_tokens=512, repetition_penalty=1.1)
 
-    return responses
+    return [response[0]["generated_text"][-1]["content"] for response in responses]
 
 def classify_pairs(model, tokenizer, prompts, responses):
     texts = [prompt + "[SEP]" + response for prompt, response in zip(prompts, responses)]
@@ -109,7 +88,7 @@ def generate(submission: list[dict[str, str]], team_id: str) -> list[dict[str, s
     ids = [s["id"] for s in submission]
     prompts = [s["prompt"] for s in submission]
 
-    responses = generate_responses(chat_model, chat_tokenizer, prompts)
+    responses = generate_responses(pipeline, chat_tokenizer, prompts)
     print(responses)
 
     scores = classify_pairs(cls_model, cls_tokenizer, prompts, responses)