Update tasks/text.py

tasks/text.py

@@ -6,17 +6,17 @@ import random
 
 from .utils.evaluation import TextEvaluationRequest
 from .utils.emissions import tracker, clean_emissions_data, get_space_info
+from .utils.preprocessing import process_text,predict 
+
 
 #packages needed for inference
-from sentence_transformers import SentenceTransformer
-from xgboost import XGBClassifier
 import pickle
 import torch
 import os
 
 router = APIRouter()
 
-DESCRIPTION = "Embedding + Neural Network"
+DESCRIPTION = "TF-IDF + RF"
 ROUTE = "/text"
 
 @router.post(ROUTE, tags=["Text Task"], 
@@ -62,59 +62,17 @@ async def evaluate_text(request: TextEvaluationRequest):
     #--------------------------------------------------------------------------------------------
     # YOUR MODEL INFERENCE CODE HERE
 
-    # Set the device to MPS (if available)
-    device = torch.device("mps" if torch.backends.mps.is_available() else "cpu")
-    print(f"Using device: {device}")
-
-    model_name = "sentence-transformers/paraphrase-multilingual-mpnet-base-v2"  # You can use other Sentence Transformers models as needed
-    sentence_model = SentenceTransformer(model_name)
 
-    # Convert each sentence into a vector representation (embedding)
-    embeddings = sentence_model.encode(test_dataset['quote'], convert_to_tensor=True)
 
     # Make random predictions (placeholder for actual model inference)
     true_labels = test_dataset["label"]
 
-    """
-    from torch import nn, optim
-
-    class SimpleNN2(nn.Module):
-        def __init__(self, input_dim, output_dim):
-            super(SimpleNN2, self).__init__()
-            self.fc1 = nn.Linear(input_dim, 128)  # Reduce hidden units
-            self.fc2 = nn.Linear(128, 64)  # Further reduce units
-            self.fc3 = nn.Linear(64, output_dim)
-            self.relu = nn.ReLU()
-            self.dropout = nn.Dropout(0.3)  # Add dropout
-            self.batch_norm1 = nn.BatchNorm1d(128)
-            self.batch_norm2 = nn.BatchNorm1d(64)
-
-        def forward(self, x):
-            x = self.relu(self.batch_norm1(self.fc1(x)))
-            x = self.dropout(x)  # Apply dropout
-            x = self.relu(self.batch_norm2(self.fc2(x)))
-            x = self.dropout(x)  # Apply dropout
-            x = self.fc3(x)  # Output raw logits
-            return x
-    """
-
     current_file_path = os.path.abspath(__file__)
     current_dir = os.path.dirname(current_file_path)
 
-    # model_nn = torch.load(os.path.join(current_dir,"model_nn.pth"), map_location=device)
-    model_nn = torch.jit.load(os.path.join(current_dir,"model_nn_scripted.pth"), map_location=device)  
-        
-
-    # Set the model to evaluation mode
-    model_nn.eval()
-
-    # Make predictions
-    with torch.no_grad():
-        outputs = model_nn(embeddings)
-        _, predicted = torch.max(outputs, 1)  # Get the class with the highest score
-
-    # Decode the predictions back to original labels using label_encoder
-    predictions = predicted.cpu().numpy()
+    # Make predictions using the loaded model
+    predictions = predict(test_dataset, os.path.join(current_dir,"tf-idf_vectorizer.pkl") ,os.path.join(current_dir,"random_forest_model.pkl"))
+    predictions = [LABEL_MAPPING[label] for label in predictions]
 
     #--------------------------------------------------------------------------------------------
     # YOUR MODEL INFERENCE STOPS HERE