Add overview of magpie battle

.gitignore

@@ -0,0 +1 @@
+.env

README.md

@@ -1,6 +1,6 @@
 ---
-title: Data Viber
-emoji: 📉
+title: Dataset Viber - chat preference magpie battle
+emoji: ⚔️
 colorFrom: red
 colorTo: yellow
 sdk: gradio

app.py

@@ -1,54 +1,42 @@
 import os
-import io
 import random
 
-import requests
-from PIL import Image
 from dataset_viber import AnnotatorInterFace
-
-HF_TOKEN = os.environ["HF_TOKEN"]
-HEADERS = {"Authorization": f"Bearer {HF_TOKEN}"}
-DATASET_SERVER_URL = "https://datasets-server.huggingface.co"
-DATASET_NAME = "poloclub%2Fdiffusiondb&config=2m_random_1k&split=train"
-MODEL_URL = (
-    "https://api-inference.huggingface.co/models/black-forest-labs/FLUX.1-schnell"
-)
-
-
-def retrieve_sample(idx):
-    api_url = f"{DATASET_SERVER_URL}/rows?dataset={DATASET_NAME}&offset={idx}&length=1"
-    response = requests.get(api_url, headers=HEADERS)
-    data = response.json()
-    img_url = data["rows"][0]["row"]["image"]["src"]
-    prompt = data["rows"][0]["row"]["prompt"]
-    return img_url, prompt
-
-
-def get_rows():
-    api_url = f"{DATASET_SERVER_URL}/size?dataset={DATASET_NAME}"
-    response = requests.get(api_url, headers=HEADERS)
-    num_rows = response.json()["size"]["config"]["num_rows"]
-    return num_rows
-
-
-def generate_response(prompt):
-    payload = {
-        "inputs": prompt,
-    }
-    response = requests.post(MODEL_URL, headers=HEADERS, json=payload)
-    image = Image.open(io.BytesIO(response.content))
-    return image
-
+from datasets import load_dataset
+from huggingface_hub import InferenceClient
+
+MODEL_IDS = [
+    "meta-llama/Meta-Llama-3.1-8B-Instruct",
+    "microsoft/Phi-3-mini-4k-instruct",
+    "mistralai/Mistral-7B-Instruct-v0.2"
+]
+CLIENTS = [InferenceClient(model_id, token=os.environ["HF_AUTH_TOKEN_PERSONAL"]) for model_id in MODEL_IDS]
+
+dataset = load_dataset("argilla/magpie-ultra-v0.1", split="train")
+
+def _get_response(messages):
+    client = random.choice(CLIENTS)
+    message = client.chat_completion(
+        messages=messages,
+        stream=False,
+        max_tokens=2000
+    )
+    return message.choices[0].message.content
 
 def next_input(_prompt, _completion_a, _completion_b):
-    random_idx = random.randint(0, get_rows()) - 1
-    img_url, prompt = retrieve_sample(random_idx)
-    generated_image = generate_response(prompt)
-    return (prompt, img_url, generated_image)
+    new_dataset = dataset.shuffle()
+    row = new_dataset[0]
+    messages = row["messages"][:-1]
+    completions = [row["response"]]
+    completions.append(_get_response(messages))
+    completions.append(_get_response(messages))
+    random.shuffle(completions)
+    return messages, completions.pop(), completions.pop()
+
 
 if __name__ == "__main__":
-    interface = AnnotatorInterFace.for_image_generation_preference(
-        fn=next_input,
+    interface = AnnotatorInterFace.for_chat_generation_preference(
+        fn_next_input=next_input,
         dataset_name=None,
     )
     interface.launch()

requirements.txt

@@ -1 +1 @@
-git+https://github.com/burtenshaw/data-viber.git@example/image-generation-preference#egg=dataset_viber
+dataset-viber==0.2.1