Update main.py

main.py

@@ -311,6 +311,42 @@ async def fast_classify_text(statement: str = Form(...)):
     except Exception as e:
         # Handle general errors
         return JSONResponse(content=f"Error in classification pipeline: {str(e)}", status_code=500)
+
+app.post("/fast_classify_v2/", description="Quickly classify text into predefined categories with confidence scores.")
+async def fast_classify_text(statement: str = Form(...)):
+    try:
+        # Use run_in_executor to handle the synchronous model call asynchronously
+        loop = asyncio.get_running_loop()
+        result = await loop.run_in_executor(
+            executor, 
+            lambda: nlp_sequence_classification(statement, labels, multi_label=False)
+        )
+
+        # Extract all labels and their scores
+        all_labels = result["labels"]
+        all_scores = result["scores"]
+
+        # Extract the best label and score
+        best_label = all_labels[0]
+        best_score = all_scores[0]
+
+        # Prepare the response
+        full_response = {
+            "classification": best_label,
+            "confidence": best_score,
+            "all_labels": {label: score for label, score in zip(all_labels, all_scores)}
+        }
+
+        return full_response
+    except asyncio.TimeoutError:
+        # Handle timeout
+        return JSONResponse(content="Classification timed out. Try a shorter input or increase timeout.", status_code=504)
+    except HTTPException as http_exc:
+        # Handle HTTP errors
+        return JSONResponse(content=f"HTTP error: {http_exc.detail}", status_code=http_exc.status_code)
+    except Exception as e:
+        # Handle general errors
+        return JSONResponse(content=f"Error in classification pipeline: {str(e)}", status_code=500)
         
 # Set up CORS middleware
 origins = ["*"]  # or specify your list of allowed origins