add exception catch

app.py

@@ -27,31 +27,35 @@ def predict(
         instruction += f"<|im_start|>user\n{user}\n<|im_end|>\n<|im_start|>assistant\n{assistant}\n<|im_end|>\n"
 
     instruction += f"<|im_start|>user\n{message}\n<|im_end|>\n<|im_start|>assistant\n"
-    print(instruction)
-    streamer = TextIteratorStreamer(
-        tokenizer,
-        skip_prompt=True,
-        skip_special_tokens=True,
-    )
-    enc = tokenizer(instruction, return_tensors="pt", padding=True, truncation=True)
-    input_ids, attention_mask = enc.input_ids, enc.attention_mask
-    if input_ids.shape[1] > CONTEXT_LENGTH:
-        input_ids = input_ids[:, -CONTEXT_LENGTH:]
-        attention_mask = attention_mask[:, -CONTEXT_LENGTH:]
+    try:
+        streamer = TextIteratorStreamer(
+            tokenizer,
+            skip_prompt=True,
+            skip_special_tokens=True,
+        )
+        enc = tokenizer(instruction, return_tensors="pt", padding=True, truncation=True)
+        input_ids, attention_mask = enc.input_ids, enc.attention_mask
+        if input_ids.shape[1] > CONTEXT_LENGTH:
+            input_ids = input_ids[:, -CONTEXT_LENGTH:]
+            attention_mask = attention_mask[:, -CONTEXT_LENGTH:]
+
+        generate_kwargs = dict(
+            input_ids=input_ids.to(device),
+            attention_mask=attention_mask.to(device),
+            streamer=streamer,
+            do_sample=True,
+            temperature=temperature,
+            max_new_tokens=max_new_tokens,
+            top_k=top_k,
+            repetition_penalty=repetition_penalty,
+            top_p=top_p,
+        )
+        t = Thread(target=model.generate, kwargs=generate_kwargs)
+        t.start()
+
+    except Exception as e:
+        streamer = f"{e}"
 
-    generate_kwargs = dict(
-        input_ids=input_ids.to(device),
-        attention_mask=attention_mask.to(device),
-        streamer=streamer,
-        do_sample=True,
-        temperature=temperature,
-        max_new_tokens=max_new_tokens,
-        top_k=top_k,
-        repetition_penalty=repetition_penalty,
-        top_p=top_p,
-    )
-    t = Thread(target=model.generate, kwargs=generate_kwargs)
-    t.start()
     outputs = []
     for new_token in streamer:
         outputs.append(new_token)