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_prompt_classification = pipeline("text2text-generation", model="google/flan-t5-large")
+nlp_sequence_classification = pipeline("zero-shot-classification", model="facebook/bart-large-mnli")
 
 description = """
 ## Image-based Document QA
@@ -264,25 +264,37 @@ async def test_transcription(file: UploadFile = File(...)):
     except Exception as e:
         raise HTTPException(status_code=500, detail=f"Error during transcription: {str(e)}")
 
-@app.post("/prompt_classify/", description="Classify the provided statement into one of the predefined categories.")
-async def prompt_classify_text(statement: str = Form(...)):
+# Predefined classifications
+labels = [
+    "All Pricing copy quote requested",
+    "Change to quote",
+    "Change to quote & Status Check",
+    "Change to quote (Items missed?)",
+    "Confirmation",
+    "Copy quote requested",
+    "Cost copy quote requested",
+    "MRSP copy quote requested",
+    "MSRP & All Pricing copy quote requested",
+    "MSRP & Cost copy quote requested",
+    "No narrative in email",
+    "Notes not clear",
+    "Retail copy quote requested",
+    "Status Check (possibly)"
+]
+
+@app.post("/fast_classify/", description="Quickly classify text into predefined categories.")
+async def fast_classify_text(statement: str = Form(...)):
     try:
-        # Predefined prompt with placeholders
-        prompt = (
-            "Please classify the statement in one of the following classifications: "
-            "Negative, Neutral, Positive"
-            f"Statement: {statement}"
-        )
+        # Use zero-shot classification to classify statement into one of the provided labels
+        result = nlp_sequence_classification(statement, labels, multi_label=False)
 
-        # Generate the response based on the prompt
-        result = nlp_prompt_classification(prompt, max_length=50, num_return_sequences=1)
+        # Extract the best label and score
+        best_label = result["labels"][0]
+        best_score = result["scores"][0]
 
-        # Extract the generated classification from the response
-        classification = result[0]['generated_text'].strip()
-
-        return {"classification": classification}
+        return {"classification": best_label, "confidence": best_score}
     except Exception as e:
-        return JSONResponse(content=f"Error in prompt classification: {str(e)}", status_code=500)
+        return JSONResponse(content=f"Error in classification: {str(e)}", status_code=500)
         
 # Set up CORS middleware
 origins = ["*"]  # or specify your list of allowed origins