Update app.py

app.py

@@ -133,17 +133,15 @@ def label_single_response_with_model(model_id, story, question, criteria, respon
         predicted_class = torch.argmax(logits, dim=1).item()
         return str(predicted_class)
     else:
+        # Use structured JSON generation like in the original script
         model = get_outlines_model(model_id, DEVICE_MAP, QUANTIZATION_BITS)
-        generator = Generator(model)
-        with torch.no_grad():
-            result = generator(prompt)
+        sampler = outlines.samplers.greedy()  # Match original temperature=0 behavior
+        generator = outlines.generate.json(model, ResponseModel, sampler=sampler)
+        result = generator(prompt)
         return result.score
 
-
 @spaces.GPU
-def label_multi_responses_with_model(
-    model_id, story, question, criteria, response_file
-):
+def label_multi_responses_with_model(model_id, story, question, criteria, response_file):
     df = pd.read_csv(response_file.name)
     assert "response" in df.columns, "CSV must contain a 'response' column."
     prompts = [
@@ -158,10 +156,11 @@ def label_multi_responses_with_model(
         predicted_classes = torch.argmax(logits, dim=1).tolist()
         scores = [str(cls) for cls in predicted_classes]
     else:
+        # Use structured JSON generation for batch processing
         model = get_outlines_model(model_id, DEVICE_MAP, QUANTIZATION_BITS)
-        generator = Generator(model)
-        with torch.no_grad():
-            results = generator(prompts)
+        sampler = outlines.samplers.greedy()
+        generator = outlines.generate.json(model, ResponseModel, sampler=sampler)
+        results = generator(prompts)
         scores = [r.score for r in results]
 
     df["score"] = scores