AroojImtiaz commited on
Commit
494cdfc
·
verified ·
1 Parent(s): 2d0a1cf

Create main.py

Browse files
Files changed (1) hide show
  1. main.py +38 -0
main.py ADDED
@@ -0,0 +1,38 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import re
2
+ import string
3
+ import nltk
4
+ from fastapi import FastAPI, HTTPException, RedirectResponse
5
+ from pydantic import BaseModel
6
+ from transformers import pipeline
7
+
8
+ # Initialize FastAPI app and download necessary NLTK resources
9
+ app = FastAPI()
10
+ nltk.download('punkt')
11
+ nltk.download('wordnet')
12
+
13
+ # Text preprocessing and model loading
14
+ def preprocess_text(text):
15
+ text = re.sub(r'http[s]?://\S+', '', text) # Remove URLs
16
+ text = re.sub(r'[' + re.escape(string.punctuation) + ']', '', text) # Remove punctuation
17
+ text = text.lower() # Convert to lowercase
18
+ lemmatizer = nltk.WordNetLemmatizer()
19
+ return ' '.join([lemmatizer.lemmatize(w) for w in nltk.word_tokenize(text)]) # Lemmatize
20
+
21
+ model = pipeline("text-classification", model="j-hartmann/emotion-english-distilroberta-base")
22
+
23
+ # API endpoints
24
+ @app.get('/')
25
+ async def redirect_to_docs():
26
+ return RedirectResponse(url='/docs')
27
+
28
+ @app.post('/analyze/')
29
+ async def predict_sentiment(text: TextInput):
30
+ try:
31
+ processed_text = preprocess_text(text.text)
32
+ return model(processed_text)
33
+ except Exception as e:
34
+ raise HTTPException(status_code=500, detail=str(e))
35
+
36
+ # Data model
37
+ class TextInput(BaseModel):
38
+ text: str