Update app.py

app.py

@@ -2,6 +2,7 @@ import gradio as gr
 import requests
 import os  
 import json
+import time
 import transformers
 from transformers import AutoTokenizer, AutoModelForCausalLM
 
@@ -30,11 +31,33 @@ def check_hallucination(assertion, citation):
     header = {"Authorization": f"Bearer {hf_token}"}
     payload = {"inputs": f"{assertion} [SEP] {citation}"}
 
-    response = requests.post(api_url, headers=header, json=payload, timeout=120)
-    output = response.json()
-    output = output[0][0]["score"]
+    attempts = 0
+    max_attempts = 3
+    wait_time = 180  # 3 minutes
+
+    while attempts < max_attempts:
+        try:
+            response = requests.post(api_url, headers=header, json=payload, timeout=120)
+            response.raise_for_status()  # This will raise an exception for HTTP error codes
+            output = response.json()
+            output = output[0][0]["score"]
+            return f"**hallucination score:** {output}"
+        except requests.exceptions.HTTPError as http_err:
+            print(f"HTTP error occurred: {http_err}")  # Python 3.6
+        except requests.exceptions.RequestException as err:
+            print(f"Other error occurred: {err}")  # Python 3.6
+        except KeyError:
+            print("KeyError: The expected key was not found in the response. The endpoint might be waking up.")
+        
+        attempts += 1
+        if attempts < max_attempts:
+            print(f"Attempt {attempts} failed. Waiting for {wait_time} seconds before retrying...")
+            time.sleep(wait_time)
+        else:
+            print("Maximum attempts reached. Please try again later.")
+            return "Error: Unable to retrieve hallucination score after multiple attempts."
 
-    return f"**hallucination score:** {output}"
+    return "Error: Unable to process the hallucination check."
 
 
 def query_vectara(text):