Grinding's picture
Update app.py
30b3ef6 verified
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from transformers import pipeline
from typing import List
app = FastAPI()
class SentenceListPayload(BaseModel):
sentences: List[str]
# Load the text classification pipeline on startup
try:
action_item_classifier = pipeline(
"text-classification",
model="knkarthick/Action_Items",
device="cpu",
)
print("✅ Action item model loaded successfully")
except Exception as e:
action_item_classifier = None
print(f"❌ Error loading action item model: {e}")
@app.post("/classify-action-items")
async def classify_sentences(payload: SentenceListPayload):
if not action_item_classifier:
raise HTTPException(status_code=503, detail="Action item model is not available.")
results = action_item_classifier(payload.sentences)
# Filter for sentences classified as action items with a confidence threshold
action_items = []
for i, sentence in enumerate(payload.sentences):
if results[i]['label'] == 'LABEL_1' and results[i]['score'] > 0.8:
action_items.append({
"sentence": sentence,
"confidence": results[i]['score']
})
return {"action_items": action_items}