Update main.py

main.py

@@ -12,6 +12,20 @@ import json
 import pandas as pd
 from datetime import datetime
 
+import spacy
+from simpletransformers.ner import NERModel
+import json
+import fasttext
+
+labels_file = f"ml_models/labels.json"
+ner_model_directory = f"ml_models/ner_model/"
+sentiment_model_file = f"ml_models/sentiment_model/model.ft"
+
+
+LANGUAGE_MODEL = spacy.load('en_core_web_sm')
+LABELS = ['O', 'B-AMENITIES', 'I-AMENITIES', 'I-CLEANLINESS', 'B-CLEANLINESS', 'I-COMMUNICATION', 'B-COMMUNICATION', 'B-CONDITION', 'I-CONDITION', 'I-CUSTOMER_SERVICE', 'B-CUSTOMER_SERVICE', 'B-EXTERIOR_LIGHTING', 'I-EXTERIOR_LIGHTING', 'B-FINANCIAL', 'I-FINANCIAL', 'B-INTERIOR_LIGHTING', 'I-INTERIOR_LIGHTING', 'B-INTERNET', 'I-INTERNET', 'B-LANDSCAPING_GROUNDS', 'I-LANDSCAPING_GROUNDS', 'B-MAINTENANCE_CLEANLINESS', 'I-MAINTENANCE_CLEANLINESS', 'I-MAINTENANCE_SERVICE', 'B-MAINTENANCE_SERVICE', 'B-MAINTENANCE_TIMELINESS', 'I-MAINTENANCE_TIMELINESS', 'B-MOVE_IN_QUALITY', 'I-MOVE_IN_QUALITY', 'I-NOISE', 'B-NOISE', 'B-PACKAGES_MAIL', 'I-PACKAGES_MAIL', 'B-PARKING', 'I-PARKING', 'I-PESTS', 'B-PESTS', 'B-PET_WASTE', 'I-PET_WASTE', 'I-SECURITY', 'B-SECURITY', 'B-SMOKE', 'I-SMOKE', 'B-TRASH', 'I-TRASH']
+SENTIMENT_MODEL = fasttext.load_model(sentiment_model_file)
+
 
 app = FastAPI()
 
@@ -39,7 +53,7 @@ def predict_single_review(text_rating: TextRatingRequest):
     #     spans = [{'label': text, 'color': '', 'value': '', 'sentiment': '', 'score': ''}]
     #     has_sentiments = False
     # processed_data['spans'] = spans
-    processed_data, has_sentiments = process_single_comment(raw_data)
+    processed_data, has_sentiments = process_single_comment(raw_data, LANGUAGE_MODEL, SENTIMENT_MODEL, LABELS )
     end_time = time.time()
     print(f"Time taken to process the data : {end_time - start_time} seconds")
     return {"processed_data": processed_data, "has_sentiments":has_sentiments}
@@ -52,7 +66,7 @@ def predict_all_reviews(reviews: AllTextRatingRequest):
     start_time = time.time()
     for review in reviews:
         raw_data = {"text":review.get('text', str()), "star_rating":review.get('rating', 5), "skip":skip}
-        processed_data, has_sentiments = process_single_comment(raw_data)
+        processed_data, has_sentiments = process_single_comment(raw_data, LANGUAGE_MODEL, SENTIMENT_MODEL, LABELS )
         processed_data_list.append({"processed_data":processed_data, "has_sentiments":has_sentiments})
     end_time = time.time()
     print(f"Time taken to process the data : {end_time - start_time} seconds")
@@ -75,7 +89,7 @@ def predict_file_responses(file: UploadFile = File(...)):
             star_rating = row['STAR RATING']
             review_id = row['REVIEWID']
             raw_data = {"text":text, "star_rating":star_rating, "skip":False}
-            processed_data, has_sentiments = process_single_comment(raw_data) 
+            processed_data, has_sentiments = process_single_comment(raw_data,LANGUAGE_MODEL, SENTIMENT_MODEL, LABELS ) 
             now = datetime.now()
             print(f"Processed review with index {index} at time {now.time()}")
             processed_data_list.append({"processed_data":processed_data, "has_sentiments":has_sentiments, "review_id":review_id})