Update main.py

main.py

@@ -1,154 +1,76 @@
-from contextlib import asynccontextmanager
-from fastapi import FastAPI, HTTPException
-from pydantic import BaseModel, ValidationError
-from fastapi.encoders import jsonable_encoder
-
-# TEXT PREPROCESSING
-# --------------------------------------------------------------------
 import re
 import string
 import nltk
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel
+from typing import Optional
+from transformers import pipeline
+from pyngrok import ngrok
+import nest_asyncio
+from fastapi.responses import RedirectResponse
+
+# Download NLTK resources
 nltk.download('punkt')
 nltk.download('wordnet')
-nltk.download('omw-1.4')
-from nltk.stem import WordNetLemmatizer
 
-# Function to remove URLs from text
+# Initialize FastAPI app
+app = FastAPI()
+
+# Text preprocessing functions
 def remove_urls(text):
     return re.sub(r'http[s]?://\S+', '', text)
 
-# Function to remove punctuations from text
 def remove_punctuation(text):
     regular_punct = string.punctuation
-    return str(re.sub(r'['+regular_punct+']', '', str(text)))
+    return re.sub(r'['+regular_punct+']', '', text)
 
-# Function to convert the text into lower case
 def lower_case(text):
     return text.lower()
 
-# Function to lemmatize text
 def lemmatize(text):
-    wordnet_lemmatizer = WordNetLemmatizer()
-
+    wordnet_lemmatizer = nltk.WordNetLemmatizer()
     tokens = nltk.word_tokenize(text)
-    lemma_txt = ''
-    for w in tokens:
-        lemma_txt = lemma_txt + wordnet_lemmatizer.lemmatize(w) + ' '
+    return ' '.join([wordnet_lemmatizer.lemmatize(w) for w in tokens])
 
-    return lemma_txt
+# Model loading
+lyx_pipe = pipeline("text-classification", model="lxyuan/distilbert-base-multilingual-cased-sentiments-student")
 
-def preprocess_text(text):
-    # Preprocess the input text
-    text = remove_urls(text)
-    text = remove_punctuation(text)
-    text = lower_case(text)
-    text = lemmatize(text)
-    return text
-
-# Load the model using FastAPI lifespan event so that the model is loaded at the beginning for efficiency
-@asynccontextmanager
-async def lifespan(app: FastAPI):
-    # Load the model from HuggingFace transformers library
-    from transformers import pipeline
-    global sentiment_task
-    sentiment_task = pipeline("sentiment-analysis", model="cardiffnlp/twitter-roberta-base-sentiment-latest", tokenizer="cardiffnlp/twitter-roberta-base-sentiment-latest")
-    yield
-    # Clean up the model and release the resources
-    del sentiment_task
-
-# Initialize the FastAPI app
-app = FastAPI(lifespan=lifespan)
-
-# Define the input data model
+# Input data model
 class TextInput(BaseModel):
     text: str
 
-# Define the welcome endpoint
+# Welcome endpoint
 @app.get('/')
 async def welcome():
-    return "Welcome to our Text Classification API"
+    # Redirect to the Swagger UI page
+    return RedirectResponse(url="/docs")
 
-# Validate input text length
-MAX_TEXT_LENGTH = 1000
+# Sentiment analysis endpoint
+@app.post('/analyze/')
+async def Predict_Sentiment(text_input: TextInput):    
+    text = text_input.text
 
-# Define the sentiment analysis endpoint 
-@app.post('/analyze/{text}')
-async def classify_text(text_input:TextInput):    
-    try:
-        # Convert input data to JSON serializable dictionary
-        text_input_dict = jsonable_encoder(text_input)
-        # Validate input data using Pydantic model
-        text_data = TextInput(**text_input_dict)  # Convert to Pydantic model
-
-        # Validate input text length
-        if len(text_input.text) > MAX_TEXT_LENGTH:
-            raise HTTPException(status_code=400, detail="Text length exceeds maximum allowed length")
-        elif len(text_input.text) == 0:
-            raise HTTPException(status_code=400, detail="Text cannot be empty")
-    except ValidationError as e:
-        # Handle validation error
-        raise HTTPException(status_code=422, detail=str(e))
+    # Text preprocessing
+    text = remove_urls(text)
+    text = remove_punctuation(text)
+    text = lower_case(text)
+    text = lemmatize(text)
 
+    # Perform sentiment analysis
     try:
-        # Perform text classification
-        return sentiment_task(preprocess_text(text_input.text))
-    except ValueError as ve:
-        # Handle value error
-        raise HTTPException(status_code=400, detail=str(ve))
+        return lyx_pipe(text)
     except Exception as e:
-        # Handle other server errors
         raise HTTPException(status_code=500, detail=str(e))
 
-# Load the model using FastAPI lifespan event so that the model is loaded at the beginning for efficiency
-@asynccontextmanager
-async def lifespan(app: FastAPI):
-    # Load the model from HuggingFace transformers library
-    from transformers import pipeline
-    global sentiment_task
-    sentiment_task = pipeline("sentiment-analysis", model="cardiffnlp/twitter-roberta-base-sentiment-latest", tokenizer="cardiffnlp/twitter-roberta-base-sentiment-latest")
-    yield
-    # Clean up the model and release the resources
-    del sentiment_task
-
-# Initialize the FastAPI app
-app = FastAPI(lifespan=lifespan)
+# Run the FastAPI app using Uvicorn
+if __name__ == "__main__":
+    # Create ngrok tunnel
+    ngrok_tunnel = ngrok.connect(7860)
+    print('Public URL:', ngrok_tunnel.public_url)
 
-# Define the input data model
-class TextInput(BaseModel):
-    text: str
+    # Allow nested asyncio calls
+    nest_asyncio.apply()
 
-# Define the welcome endpoint
-@app.get('/')
-async def welcome():
-    return "Welcome to our Text Classification API"
-
-# Validate input text length
-MAX_TEXT_LENGTH = 1000
-
-# Define the sentiment analysis endpoint 
-@app.post('/analyze/{text}')
-async def classify_text(text_input:TextInput):    
-    try:
-        # Convert input data to JSON serializable dictionary
-        text_input_dict = jsonable_encoder(text_input)
-        # Validate input data using Pydantic model
-        text_data = TextInput(**text_input_dict)  # Convert to Pydantic model
-
-        # Validate input text length
-        if len(text_input.text) > MAX_TEXT_LENGTH:
-            raise HTTPException(status_code=400, detail="Text length exceeds maximum allowed length")
-        elif len(text_input.text) == 0:
-            raise HTTPException(status_code=400, detail="Text cannot be empty")
-    except ValidationError as e:
-        # Handle validation error
-        raise HTTPException(status_code=422, detail=str(e))
-
-    try:
-        # Perform text classification
-        return sentiment_task(preprocess_text(text_input.text))
-    except ValueError as ve:
-        # Handle value error
-        raise HTTPException(status_code=400, detail=str(ve))
-    except Exception as e:
-        # Handle other server errors
-        raise HTTPException(status_code=500, detail=str(e))
+    # Run the FastAPI app with Uvicorn
+    import uvicorn
+    uvicorn.run(app, port=7860)