Update main.py

main.py

@@ -33,29 +33,34 @@ def generate_text(
     body: dict = Body(...),
     api_key: str = Depends(get_api_key)
 ):
-    print("[request]",request)
-    print("[body]",body)
-    prompt = body.get("prompt", "")
-    system_prompt = body.get("sysPrompt", "")
-    temperature = body.get("temperature", 0.5)
-    top_p = body.get("top_p", 0.95)
-    max_new_tokens = body.get("max_new_tokens",512)
-    repetition_penalty = body.get("repetition_penalty", 1.0)
-    print(f"temperature + {temperature}")
-    history = []  # You might need to handle this based on your actual usage
-    formatted_prompt = format_prompt(f"{system_prompt}, {prompt}", history)
-
-    stream = client.text_generation(
-        formatted_prompt,
-        temperature=temperature,
-        max_new_tokens=max_new_tokens,
-        top_p=top_p,
-        repetition_penalty=repetition_penalty,
-        do_sample=True,
-        seed=random.randint(0, 10**7),
-        stream=False,
-        details=False,
-        return_full_text=False
-    )
-
-    return {"generated_text": stream}
+    try:
+        print("[request]",request)
+        print("[body]",body)
+        prompt = body.get("prompt", "")
+        system_prompt = body.get("sysPrompt", "")
+        temperature = body.get("temperature", 0.5)
+        top_p = body.get("top_p", 0.95)
+        max_new_tokens = body.get("max_new_tokens",512)
+        repetition_penalty = body.get("repetition_penalty", 1.0)
+        print(f"temperature + {temperature}")
+        history = []  # You might need to handle this based on your actual usage
+        formatted_prompt = format_prompt(f"{system_prompt}, {prompt}", history)
+
+        stream = client.text_generation(
+            formatted_prompt,
+            temperature=temperature,
+            max_new_tokens=max_new_tokens,
+            top_p=top_p,
+            repetition_penalty=repetition_penalty,
+            do_sample=True,
+            seed=random.randint(0, 10**7),
+            stream=False,
+            details=False,
+            return_full_text=False
+        )
+
+        return {"generated_text": stream}
+    except error:
+        print("[error Request]", request)
+        return {error}
+