Update main.py

main.py

@@ -267,7 +267,10 @@ async def test_transcription(file: UploadFile = File(...)):
     except Exception as e:
         raise HTTPException(status_code=500, detail=f"Error during transcription: {str(e)}")
 
-# Predefined classification labels
+# Define the ThreadPoolExecutor globally to manage asynchronous execution
+executor = ThreadPoolExecutor(max_workers=10)
+
+# Predefined classifications
 labels = [
     "All Pricing copy quote requested",
     "Change to quote",
@@ -285,58 +288,30 @@ labels = [
     "Status Check (possibly)"
 ]
 
-# Keyword map for rule-based shortcuts
-keyword_map = {
-    "All Pricing copy quote requested": ["Per ATP", "Not Want", "No Advertising"],
-    "Change to quote": ["MSRP"],
-    "Change to quote & Status Check": ["Send clean Quote"],
-    "Change to quote (Items missed?)": ["Per ATP"],
-    "Confirmation": ["N/A"],
-    "Copy quote requested": ["N/A"],
-    "Cost copy quote requested": ["N/A"],
-    "MRSP copy quote requested": ["Quote Requested", "Revise Quote", "Include the following items"],
-    "MSRP & All Pricing copy quote requested": ["Revise and resend Quote"],
-    "MSRP & Cost copy quote requested": ["Per RTS"],
-    "No narrative in email": ["Requesting"],
-    "Notes not clear": ["Add to the quote"],
-    "Retail copy quote requested": ["Action Required", "Send copy of Quote", "All Pricing"],
-    "Status Check (possibly)": ["Re-quote", "Current date and Pricing", "Send quote"]
-}
-
-# Thread pool executor for handling requests concurrently
-executor = ThreadPoolExecutor(max_workers=10)
-
-# Helper function to check keywords
-def check_keywords(statement):
-    for label, keywords in keyword_map.items():
-        for keyword in keywords:
-            if re.search(rf"\b{keyword}\b", statement, re.IGNORECASE):
-                return label
-    return None
-
-@app.post("/fast_classify/", description="Quickly classify text with a rule-based approach combined with model inference.")
+@app.post("/fast_classify/", description="Quickly classify text into predefined categories.")
 async def fast_classify_text(statement: str = Form(...)):
     try:
-        # Step 1: Check for rule-based keyword match
-        rule_based_classification = check_keywords(statement)
-        if rule_based_classification:
-            return {"classification": rule_based_classification, "confidence": 1.0}
-
-        # Step 2: Run model classification if no rule-based match
+        # Use run_in_executor to handle the synchronous model call asynchronously
         loop = asyncio.get_running_loop()
         result = await loop.run_in_executor(
-            executor,
+            executor, 
             lambda: nlp_sequence_classification(statement, labels, multi_label=False)
         )
 
+        # Extract the best label and score
         best_label = result["labels"][0]
         best_score = result["scores"][0]
 
         return {"classification": best_label, "confidence": best_score}
     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:
-        return JSONResponse(content=f"Error in classification: {str(e)}", status_code=500)
+        # 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