Update main.py

main.py

@@ -23,6 +23,7 @@ app.add_middleware(
 nlp_qa = pipeline("document-question-answering", model="jinhybr/OCR-DocVQA-Donut")
 nlp_qa_v2 = pipeline("document-question-answering", model="faisalraza/layoutlm-invoices")
 nlp_qa_v3 = pipeline("question-answering", model="deepset/roberta-base-squad2")
+nlp_classification = pipeline("text-classification", model="distilbert-base-uncased-finetuned-sst-2-english")
 
 description = """
 ## Image-based Document QA
@@ -116,6 +117,18 @@ async def perform_document_qa(
     except Exception as e:
         return JSONResponse(content=f"Error processing file: {str(e)}", status_code=500)
 
+@app.post("/classify/", description="Classify the provided text.")
+async def classify_text(text: str = Form(...)):
+    try:
+        # Perform text classification using the pipeline
+        result = nlp_classification(text)
+
+        # Return the classification result
+        return result
+    except Exception as e:
+        return JSONResponse(content=f"Error classifying text: {str(e)}", status_code=500)
+
+
 # Set up CORS middleware
 origins = ["*"]  # or specify your list of allowed origins
 app.add_middleware(