Update app.py

app.py

@@ -4,24 +4,40 @@ import random
 from dataset_viber import AnnotatorInterFace
 from datasets import load_dataset
 from huggingface_hub import InferenceClient
+import time
 
 # https://huggingface.co/models?inference=warm&pipeline_tag=text-generation&sort=trending
 MODEL_IDS = [
-    "sci-m-wang/Phi-3-mini-4k-instruct-sa-v0.1",
+    "google/gemma-2b-it",
     "microsoft/Phi-3-mini-4k-instruct"
 ]
 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")
 
-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 get_response(messages):
+    max_retries = 3
+    retry_delay = 3
+
+    for attempt in range(max_retries):
+        try:
+            client = random.choice(CLIENTS)
+            message = client.chat_completion(
+                messages=messages,
+                stream=False,
+                max_tokens=2000
+            )
+            return message.choices[0].message.content
+        except Exception as e:
+            if attempt < max_retries - 1:
+                print(f"An error occurred: {e}. Retrying in {retry_delay} seconds...")
+                time.sleep(retry_delay)
+            else:
+                print(f"Max retries reached. Last error: {e}")
+                raise
+
+    return None  # This line will only be reached if all retries fail
 
 def next_input(_prompt, _completion_a, _completion_b):
     new_dataset = dataset.shuffle()