Update app.py

app.py

@@ -9,6 +9,7 @@ import pandas as pd
 
 from huggingface_hub import Repository, upload_file
 from text_generation import Client
+from text_generation.errors import UnknownError
 from share_btn import (community_icon_html, loading_icon_html, share_btn_css,
                        share_js)
 
@@ -96,14 +97,14 @@ def get_inference_prompt(messages, model_name):
         prompt = "<|system|>\n<|endoftext|>\n"
         for message in messages :
             if message["role"] == "user" :
-                prompt += f"<|user|>\n{message['content']}<|endoftext|>\n<|assistant|>"
+                prompt += f"<|user|>\n{message['content']}<|endoftext|>\n<|assistant|>\n"
             else : # message["role"] == "assistant"
                 prompt += f"\n{message['content']}<|endoftext|>\n"
     elif model_name == "starChat-alpha" :
         prompt = "<|system|>\n<|end|>\n"
         for message in messages :
             if message["role"] == "user" :
-                prompt += f"<|user|>\n{message['content']}<|end|>\n<|assistant|>"
+                prompt += f"<|user|>\n{message['content']}<|end|>\n<|assistant|>\n"
             else : # message["role"] == "assistant"
                 prompt += f"\n{message['content']}<|end|>\n"
     else : # starCoder-gradio
@@ -180,40 +181,43 @@ def generate(
         stop_sequences=["<|end|>", "Question:"],
     )
 
-    stream = client.generate_stream(
-        prompt,
-        **generate_kwargs,
-    )
-
-    output = ""
-    for idx, response in enumerate(stream):
-        if response.token.special:
-            continue
-        output += response.token.text
-        if idx == 0:
-            history.append(" " + output)
-        else:
-            history[-1] = output
-
-        chat = [
-            (wrap_html_code(history[i].strip()), wrap_html_code(history[i + 1].strip()))
-            for i in range(0, len(history) - 1, 2)
-        ]
-
-        # chat = [(history[i].strip(), history[i + 1].strip()) for i in range(0, len(history) - 1, 2)]
-
-        yield chat, history, user_message, ""
-
-    if HF_TOKEN and do_save:
-        try:
-            now = datetime.datetime.now()
-            current_time = now.strftime("%Y-%m-%d %H:%M:%S")
-            print(f"[{current_time}] Pushing prompt and completion to the Hub")
-            save_inputs_and_outputs(now, prompt, output, generate_kwargs, model_name)
-        except Exception as e:
-            print(e)
-
-    return chat, history, user_message, ""
+    try :
+        stream = client.generate_stream(
+            prompt,
+            **generate_kwargs,
+        )
+    
+        output = ""
+        for idx, response in enumerate(stream):
+            if response.token.special:
+                continue
+            output += response.token.text
+            if idx == 0:
+                history.append(" " + output)
+            else:
+                history[-1] = output
+    
+            chat = [
+                (wrap_html_code(history[i].strip()), wrap_html_code(history[i + 1].strip()))
+                for i in range(0, len(history) - 1, 2)
+            ]
+    
+            # chat = [(history[i].strip(), history[i + 1].strip()) for i in range(0, len(history) - 1, 2)]
+    
+            yield chat, history, user_message, ""
+    
+        if HF_TOKEN and do_save:
+            try:
+                now = datetime.datetime.now()
+                current_time = now.strftime("%Y-%m-%d %H:%M:%S")
+                print(f"[{current_time}] Pushing prompt and completion to the Hub")
+                save_inputs_and_outputs(now, prompt, output, generate_kwargs, model_name)
+            except Exception as e:
+                print(e)
+    
+        return chat, history, user_message, ""
+    except UnknownError :
+        return "The model is currently loading. Please wait a few seconds and retry.", None, None, ""
 
 examples = [
     "Use the gradio library to create a calculator. It should take into account the 4 basic operations. A user should be able to enter 2 numbers, choose an operation and get the corresponding result.",