Update app.py

app.py

@@ -12,7 +12,7 @@ MODEL_IDS = [
 ]
 CLIENTS = [InferenceClient(model_id, token=os.environ["HF_TOKEN"]) for model_id in MODEL_IDS]
 
-dataset = load_dataset("argilla/magpie-ultra-v0.1", split="train")
+dataset = load_dataset("argilla/distilabel-capybara-dpo-7k-binarized", split="train")
 
 
 def get_response(messages):
@@ -41,8 +41,8 @@ def get_response(messages):
 def next_input(_prompt, _completion_a, _completion_b):
     new_dataset = dataset.shuffle()
     row = new_dataset[0]
-    messages = row["messages"][:-1]
-    completions = [row["response"]]
+    messages = row["chosen"][:-1]
+    completions = [row["chosen"][-1]["content"]]
     completions.append(get_response(messages))
     completions.append(get_response(messages))
     random.shuffle(completions)