Update app.py

app.py

@@ -3,6 +3,7 @@ import gradio as gr
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer, StoppingCriteria, StoppingCriteriaList, TextIteratorStreamer
 from threading import Thread
+import traceback
 
 model_path = 'infly/OpenCoder-8B-Instruct'
 
@@ -43,42 +44,46 @@ system_prompt = f"<|im_start|>{system_role}\n{system_prompt}<|im_end|>"
 def predict(message, history):
     # history = []
     # history_transformer_format = history + [[message, ""]]
-    stop = StopOnTokens()
-
-    # Formatting the input for the model.
-    # messages =  system_prompt + sft_end_token.join([sft_end_token.join([f"\n{sft_start_token}{user_role}\n" + item[0], f"\n{sft_start_token}{assistant_role}\n" + item[1]])
-    #                     for item in history_transformer_format])
-
-    model_messages = []
-    print(f'history: {history}')
-    for i, item in enumerate(history):
-        model_messages.append({"role": user_role, "content": item[0]})
-        model_messages.append({"role": assistant_role, "content": item[1]})
-
-    model_messages.append({"role": user_role, "content": message})
+    try:
+        stop = StopOnTokens()
     
-    print(f'model_messages: {model_messages}')
-
-    print(f'model_final_inputs: {tokenizer.apply_chat_template(model_messages, add_generation_prompt=True, tokenize=False)}', flash=True)
-    model_inputs = tokenizer.apply_chat_template(model_messages, add_generation_prompt=True, return_tensors="pt").to(device)
-    # model_inputs = tokenizer([messages], return_tensors="pt").to(device)
+        # Formatting the input for the model.
+        # messages =  system_prompt + sft_end_token.join([sft_end_token.join([f"\n{sft_start_token}{user_role}\n" + item[0], f"\n{sft_start_token}{assistant_role}\n" + item[1]])
+        #                     for item in history_transformer_format])
+    
+        model_messages = []
+        print(f'history: {history}')
+        for i, item in enumerate(history):
+            model_messages.append({"role": user_role, "content": item[0]})
+            model_messages.append({"role": assistant_role, "content": item[1]})
+    
+        model_messages.append({"role": user_role, "content": message})
+        
+        print(f'model_messages: {model_messages}')
     
-    streamer = TextIteratorStreamer(tokenizer, timeout=10., skip_prompt=True, skip_special_tokens=True)
-    generate_kwargs = dict(
-        model_inputs,
-        streamer=streamer,
-        max_new_tokens=1024,
-        do_sample=False,
-        # stopping_criteria=StoppingCriteriaList([stop])
-    )
-    t = Thread(target=model.generate, kwargs=generate_kwargs)
-    t.start()  # Starting the generation in a separate thread.
-    partial_message = ""
-    for new_token in streamer:
-        partial_message += new_token
-        if sft_end_token in partial_message:  # Breaking the loop if the stop token is generated.
-            break
-        yield partial_message
+        print(f'model_final_inputs: {tokenizer.apply_chat_template(model_messages, add_generation_prompt=True, tokenize=False)}', flush=True)
+        model_inputs = tokenizer.apply_chat_template(model_messages, add_generation_prompt=True, return_tensors="pt").to(device)
+        # model_inputs = tokenizer([messages], return_tensors="pt").to(device)
+        
+        streamer = TextIteratorStreamer(tokenizer, timeout=10., skip_prompt=True, skip_special_tokens=True)
+        generate_kwargs = dict(
+            model_inputs,
+            streamer=streamer,
+            max_new_tokens=1024,
+            do_sample=False,
+            # stopping_criteria=StoppingCriteriaList([stop])
+        )
+        t = Thread(target=model.generate, kwargs=generate_kwargs)
+        t.start()  # Starting the generation in a separate thread.
+        partial_message = ""
+        for new_token in streamer:
+            partial_message += new_token
+            if sft_end_token in partial_message:  # Breaking the loop if the stop token is generated.
+                break
+            yield partial_message
+
+    except Exception as e:
+        print(traceback.format_exc())
 
 
 css = """