Improved error handling

app.py

@@ -1,25 +1,38 @@
 import os
+import time
 
 import requests
 import streamlit as st
 
-
 API_URL = "https://api-inference.huggingface.co/models/andrek/nb2nn"
 
-def translate(text):
+
+def translate(text, wait=True):
     headers = {"Authorization": f"Bearer {os.environ['BEARER']}"}
     payload  = {
         "inputs": text,
     }
-    try:
-        response = requests.post(API_URL, headers=headers, json=payload)
-        json_response = response.json()[0]
-        if "error" in json_response:
-            return f"Model still loading: {json_response}"
-        elif "translation_text" in json_response:
-            return json_response["translation_text"]
-    except Exception as exc:
-        return f"Oops, something went wrong: {exc}"
+    response = requests.post(API_URL, headers=headers, json=payload)
+    json_response = response.json()
+    if (isinstance(json_response, dict)
+        and "error" in json_response
+        and "estimated_time" in json_response):
+        st.write(json_response)
+        if wait:
+            with st.spinner(json_response["error"]):
+                bar = st.progress(0)
+                for progress in range(json_response["estimated_time"]):
+                    bar.progress(progress / json_response["estimated_time"])
+                    time.sleep(1)
+                bar.empty()
+            return translate(text, wait=False)
+        else:
+            return "We could not load the model"
+    elif (isinstance(json_response, list)
+        and "translation_text" in json_response[0]):
+        return json_response[0]["translation_text"]
+    else:
+        return f"Oops, something went wrong: {str(json_response)}"
 
 
 st.set_page_config(