xclasscode's picture
Update main.py
ad38258 verified
raw
history blame
2.17 kB
from fastapi import Depends, FastAPI, HTTPException, status
from sqlalchemy.orm import Session
import schemas, models
from transformers import pipeline
import database
app = FastAPI()
models.database.Base.metadata.create_all(bind=database.engine)
pipe = pipeline("text-classification", model="lxyuan/distilbert-base-multilingual-cased-sentiments-student")
# Database session dependency
def get_db():
db = database.SessionLocal()
try:
yield db
finally:
db.close()
@app.post("/analyze_sentiment", status_code=status.HTTP_201_CREATED)
def create_sentiment_result(sentiment_result: schemas.SentimentResultCreate, text_input: str, db: Session = Depends(get_db)):
# Perform input validation
if not isinstance(text_input, str) or not text_input.strip():
raise HTTPException(status_code=status.HTTP_422_UNPROCESSABLE_ENTITY, detail="Invalid text input")
try:
text_content = text_input
sentiment_analysis_result = pipe(text_content)
# Create a new SentimentResult instance
new_sentiment_result = models.SentimentResult(
score=sentiment_analysis_result[0]['score'],
label = sentiment_analysis_result[0]['label'],
text_input=text_content
)
# Add the new SentimentResult to the database
db.add(new_sentiment_result)
db.commit()
db.refresh(new_sentiment_result)
return new_sentiment_result
except Exception as e:
raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"An error occurred: {str(e)}")
@app.delete("/sentiment/{id}", status_code=status.HTTP_204_NO_CONTENT)
def delete_sentiment_result(id: int, db: Session = Depends(get_db)):
# Retrieve the sentiment result and check for existence
sentiment_result_to_delete = db.query(models.SentimentResult).filter(models.SentimentResult.id == id).first()
if sentiment_result_to_delete is None:
raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail=f"SentimentResult with ID {id} not found")
# Delete the sentiment result
db.delete(sentiment_result_to_delete)
db.commit()