Update app.py

app.py

@@ -17,6 +17,11 @@ def load_model(repo_name):
 
 tokenizer, generator_conf, model = load_model(REPO_NAME)
 
+global_error = ''
+try:
+    generator = pipeline("text-generation", model=model, tokenizer=tokenizer)
+except Exception as e:
+    global_error =  f"Failed to load model: {str(e)}"
 
 def respond(
     message,
@@ -26,37 +31,44 @@ def respond(
     temperature,
     top_p,
 ):
-    messages = [{"role": "system", "content": system_message}]
-
-    for val in history:
-        if val[0]:
-            messages.append({"role": "user", "content": val[0]})
-        if val[1]:
-            messages.append({"role": "assistant", "content": val[1]})
-
-    messages.append({"role": "user", "content": message})
-
-    for message in messages:
-            role = "<|assistant|>" if message['role'] == 'assistant' else "<|user|>"
-            prompt += f"\n{role}\n{message['content']}\n<|end|>\n"        
-    # prompt += f"\n<|user|>\n{user_text}\n<|end|><|assistant|>\n"
-
-    # Generate the response
-    response_output = generator(
-        prompt,
-        generation_config=generator_conf,
-        max_new_tokens=64,
-        do_sample=True,
-        top_p=0.25,
-        repetition_penalty=1.2
-    )
-
-    generated_text = response_output[0]['generated_text']
+    result = 'none'
+    try:
+        messages = [{"role": "system", "content": system_message}]
+    
+        for val in history:
+            if val[0]:
+                messages.append({"role": "user", "content": val[0]})
+            if val[1]:
+                messages.append({"role": "assistant", "content": val[1]})
+    
+        messages.append({"role": "user", "content": message})
+    
+        for message in messages:
+                role = "<|assistant|>" if message['role'] == 'assistant' else "<|user|>"
+                prompt += f"\n{role}\n{message['content']}\n<|end|>\n"        
+        # prompt += f"\n<|user|>\n{user_text}\n<|end|><|assistant|>\n"
+    
+        # Generate the response
+        response_output = generator(
+            prompt,
+            generation_config=generator_conf,
+            max_new_tokens=64,
+            do_sample=True,
+            top_p=0.25,
+            repetition_penalty=1.2
+        )
+    
+        generated_text = response_output[0]['generated_text']
+    
+        # st.session_state.last_response = generated_text
+    
+        # Extract the assistant's response
+        result = generated_text[len(prompt):].strip()
+    except Exception as error:
+        result = str(error)
 
-    # st.session_state.last_response = generated_text
+    yield result
 
-    # Extract the assistant's response
-    yield generated_text[len(prompt):].strip()
     
     """
     for message in client.chat_completion(
@@ -79,7 +91,7 @@ For information on how to customize the ChatInterface, peruse the gradio docs: h
 demo = gr.ChatInterface(
     respond,
     additional_inputs=[
-        gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
+        gr.Textbox(value="You are a friendly Chatbot." + global_error, label="System message"),
         gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
         gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
         gr.Slider(