AroojImtiaz's picture
Create main.py
494cdfc verified
raw
history blame
1.21 kB
import re
import string
import nltk
from fastapi import FastAPI, HTTPException, RedirectResponse
from pydantic import BaseModel
from transformers import pipeline
# 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))
# Data model
class TextInput(BaseModel):
text: str