Update main.py

main.py

@@ -35,7 +35,7 @@ nlp_qa_v3 = pipeline("question-answering", model="deepset/roberta-base-squad2")
 nlp_classification = pipeline("text-classification", model="distilbert/distilbert-base-uncased-finetuned-sst-2-english")
 nlp_classification_v2 = pipeline("text-classification", model="cardiffnlp/twitter-roberta-base-sentiment-latest")
 nlp_speech_to_text = pipeline("automatic-speech-recognition", model="facebook/wav2vec2-base-960h")
-nlp_sequence_classification = pipeline("zero-shot-classification", model="facebook/bart-large-mnli")
+nlp_sequence_classification = pipeline("zero-shot-classification", model="valhalla/distilbart-mnli-12-1")
 
 description = """
 ## Image-based Document QA
@@ -285,14 +285,19 @@ labels = [
 @app.post("/fast_classify/", description="Quickly classify text into predefined categories.")
 async def fast_classify_text(statement: str = Form(...)):
     try:
-        # Use zero-shot classification to classify statement into one of the provided labels
-        result = nlp_sequence_classification(statement, labels, multi_label=False)
+        # Use asyncio to set a timeout for the classification
+        result = await asyncio.wait_for(
+            nlp_sequence_classification(statement, labels, multi_label=False),
+            timeout=5  # timeout in seconds
+        )
 
         # 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:
+        return JSONResponse(content="Classification timed out. Try a shorter input or increase timeout.", status_code=504)
     except Exception as e:
         return JSONResponse(content=f"Error in classification: {str(e)}", status_code=500)