group by model

app.py

@@ -61,9 +61,12 @@ def response(num_responses, model, correct, prompt_ids):
         responses = responses[responses['correct'].isin(correct)]
     if prompt_ids:
         responses = responses[responses['prompt'].isin(prompt_ids)]
-    if num_responses > len(responses):
-        num_responses = len(responses)
-    return responses.sample(num_responses)[['model', 'prompt', 'model_response', 'correct']]
+    # if num_responses > len(responses):
+    #     num_responses = len(responses)
+
+    # sample num_responses for each model
+    responses = responses.groupby('model').apply(lambda x: x.sample(num_responses) if num_responses < len(x) else x).reset_index()
+    return responses[['model', 'prompt', 'model_response', 'correct']]
 
 
 def barplot_for_prompt_id(prompt_ids, models):