AroojImtiaz's picture
Update main.py
d9cb7cb verified
raw
history blame contribute delete
No virus
1.24 kB
import re
import string
import nltk
from pydantic import BaseModel
from transformers import pipeline
from fastapi import FastAPI, HTTPException
from fastapi.responses import RedirectResponse
# Data model
class TextInput(BaseModel):
text: str
# Initialize FastAPI app and download necessary NLTK resources
app = FastAPI()
nltk.download('punkt')
nltk.download('wordnet')
# Text preprocessing and model loading
def preprocess_text(text):
text = re.sub(r'http[s]?://\S+', '', text) # Remove URLs
text = re.sub(r'[' + re.escape(string.punctuation) + ']', '', text) # Remove punctuation
text = text.lower() # Convert to lowercase
lemmatizer = nltk.WordNetLemmatizer()
return ' '.join([lemmatizer.lemmatize(w) for w in nltk.word_tokenize(text)]) # Lemmatize
model = pipeline("text-classification", model="j-hartmann/emotion-english-distilroberta-base")
# API endpoints
@app.get('/')
async def redirect_to_docs():
return RedirectResponse(url='/docs')
@app.post('/analyze/')
async def predict_sentiment(text: TextInput):
try:
processed_text = preprocess_text(text.text)
return model(processed_text)
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))