Update app.py

app.py

@@ -1,8 +1,7 @@
 import os
-import urllib
+import time
 import gradio as gr
-
-import requests
+from huggingface_hub import get_inference_endpoint
 
 endpoint_url = os.getenv('ENDPOINT_URL')
 personal_secret_token = os.getenv('PERSONAL_HF_TOKEN')
@@ -12,15 +11,7 @@ system_symbol = os.getenv('SYSTEM_SYMBOL')
 user_symbol = os.getenv('USER_SYMBOL')
 assistant_symbol = os.getenv('ASSISTANT_SYMBOL')
 
-headers = {
-	"Accept" : "application/json",
-    "Authorization": f"Bearer {personal_secret_token}",
-	"Content-Type": "application/json" 
-}
-
-def query(payload):
-	response = requests.post(endpoint_url, headers=headers, json=payload)
-	return response.json()
+endpoint = get_inference_endpoint(endpoint_url, token=personal_secret_token)
 
 def respond(
     message,
@@ -29,7 +20,16 @@ def respond(
     max_new_tokens,
     temperature,
     top_p,
+    progress=gr.Progress()
 ):
+    progress(0, desc="Starting")
+
+    while endpoint.status != "running":
+        progress(0.25, desc="Waking up model")
+        time.sleep(1)
+
+    progress(0.5, desc="Generating")
+    
     all_messages = [system_message]
 
     for val in history:
@@ -48,11 +48,13 @@ def respond(
         # stream=True,
     )
 
-    response = query({
+    response = endpoint.client.post({
         "inputs": turn_breaker.join(all_messages),
         "parameters": generation_kwargs
     })
 
+    progress(1, desc="Generating")
+
     return response
 
 
@@ -73,6 +75,7 @@ demo = gr.ChatInterface(
             label="Top-p (nucleus sampling)",
         ),
     ],
+    show_progress="full"
 )