Update main.py

main.py

@@ -35,6 +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")
 
 description = """
 ## Image-based Document QA
@@ -262,6 +263,26 @@ 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(...)):
+    try:
+        # Predefined prompt with placeholders
+        prompt = (
+            "Please classify the statement in one of the following classifications: "
+            "Negative, Neutral, Positive"
+            f"Statement: {statement}"
+        )
+
+        # Generate the response based on the prompt
+        result = nlp_prompt_classification(prompt, max_length=50, num_return_sequences=1)
+
+        # Extract the generated classification from the response
+        classification = result[0]['generated_text'].strip()
+
+        return {"classification": classification}
+    except Exception as e:
+        return JSONResponse(content=f"Error in prompt classification: {str(e)}", status_code=500)
         
 # Set up CORS middleware
 origins = ["*"]  # or specify your list of allowed origins