Update tasks/text.py

tasks/text.py

@@ -20,7 +20,7 @@ from .utils.emissions import tracker, clean_emissions_data, get_space_info
 
 router = APIRouter()
 
-DESCRIPTION = "First Baseline"
+DESCRIPTION = "Submission 2: SBERT+MLP"
 ROUTE = "/text"
 
 
@@ -28,19 +28,18 @@ if torch.cuda.is_available():
     device = torch.device("cuda")
 else:
     device = torch.device("cpu")
-print(device)
 
 
-MODEL = "modern-base" #mlp, ct, modern-base, modern-large, gte-base, gte-large
+MODEL = "mlp" #sk, mlp, ct, modern-base, modern-large, gte-base, gte-large 
 
-class ConspiracyClassification(
+class ConspiracyClassification768(
     nn.Module,
     PyTorchModelHubMixin, 
     # optionally, you can add metadata which gets pushed to the model card
 ):    
-    def __init__(self, num_classes):
+    def __init__(self, num_classes=8):
         super().__init__()
-        self.h1 = nn.Linear(384, 100)
+        self.h1 = nn.Linear(768, 100)
         self.h2 = nn.Linear(100, 100)
         self.h3 = nn.Linear(100, 100)
         self.h4 = nn.Linear(100, 50)
@@ -71,7 +70,7 @@ class CTBERT(
     PyTorchModelHubMixin, 
     # optionally, you can add metadata which gets pushed to the model card
 ):
-    def __init__(self, num_classes):
+    def __init__(self, num_classes=8):
         super().__init__()
         self.bert = BertForPreTraining.from_pretrained('digitalepidemiologylab/covid-twitter-bert-v2')    
         self.bert.cls.seq_relationship = nn.Linear(1024, num_classes)
@@ -87,7 +86,7 @@ class conspiracyModelBase(
     PyTorchModelHubMixin, 
     # optionally, you can add metadata which gets pushed to the model card
 ):    
-    def __init__(self, num_classes):
+    def __init__(self, num_classes=8):
         super().__init__()
         self.n_classes = num_classes
         self.bert = ModernBertForSequenceClassification.from_pretrained('answerdotai/ModernBERT-base', num_labels=num_classes)    
@@ -102,7 +101,7 @@ class conspiracyModelLarge(
     PyTorchModelHubMixin, 
     # optionally, you can add metadata which gets pushed to the model card
 ):    
-    def __init__(self, num_classes):
+    def __init__(self, num_classes=8):
         super().__init__()
         self.n_classes = num_classes
         self.bert = ModernBertForSequenceClassification.from_pretrained('answerdotai/ModernBERT-large', num_labels=num_classes)    
@@ -117,12 +116,10 @@ class gteModelLarge(
     PyTorchModelHubMixin, 
     # optionally, you can add metadata which gets pushed to the model card
 ):    
-    def __init__(self, num_classes):
+    def __init__(self, num_classes=8):
         super().__init__()
         self.n_classes = num_classes
-        #self.bert = ModernBertForSequenceClassification.from_pretrained('answerdotai/ModernBERT-large', num_labels=num_classes)    
         self.gte = AutoModel.from_pretrained('Alibaba-NLP/gte-large-en-v1.5', trust_remote_code=True)
-        #self.cls = nn.Linear(768, num_classes)
         self.cls = nn.Linear(1024, num_classes)
         
     def forward(self, input_ids, input_mask, input_type_ids):
@@ -136,20 +133,17 @@ class gteModel(
     PyTorchModelHubMixin, 
     # optionally, you can add metadata which gets pushed to the model card
 ):    
-    def __init__(self, num_classes):
+    def __init__(self, num_classes=8):
         super().__init__()
         self.n_classes = num_classes
-        #self.bert = ModernBertForSequenceClassification.from_pretrained('answerdotai/ModernBERT-large', num_labels=num_classes)    
         self.gte = AutoModel.from_pretrained('Alibaba-NLP/gte-base-en-v1.5', trust_remote_code=True)
         self.cls = nn.Linear(768, num_classes)
-        #self.cls = nn.Linear(1024, num_classes)
         
     def forward(self, input_ids, input_mask, input_type_ids):
         outputs = self.gte(input_ids = input_ids, attention_mask = input_mask, token_type_ids = input_type_ids)
         embeddings = outputs.last_hidden_state[:, 0]
         logits = self.cls(embeddings)
-        return logits
-        
+        return logits        
 
 @router.post(ROUTE, tags=["Text Task"], 
              description=DESCRIPTION)
@@ -187,20 +181,31 @@ async def evaluate_text(request: TextEvaluationRequest):
     test_dataset = dataset["test"]
 
     if MODEL =="mlp":
-        model = ConspiracyClassification.from_pretrained("ypesk/frugal-ai-mlp-baseline")
-        model = model.to(device)        
-        emb_model = SentenceTransformer("paraphrase-MiniLM-L3-v2")
+        model = ConspiracyClassification768.from_pretrained("ypesk/frugal-ai-EURECOM-mlp-768-fullset")
+        model = model.to(device)
+        emb_model = SentenceTransformer("sentence-transformers/sentence-t5-large")
         batch_size = 6
     
         test_texts = torch.Tensor(emb_model.encode([t['quote'] for t in test_dataset]))
         test_data = TensorDataset(test_texts)
         test_sampler = SequentialSampler(test_data)
         test_dataloader = DataLoader(test_data, sampler=test_sampler, batch_size=batch_size)
-        
+    elif MODEL == "sk":
+        emb_model = SentenceTransformer("sentence-transformers/sentence-t5-large")
+        batch_size = 512
+    
+        test_texts = torch.Tensor(emb_model.encode([t['quote'] for t in test_dataset]))
+        test_data = TensorDataset(test_texts)
+        test_sampler = SequentialSampler(test_data)
+        test_dataloader = DataLoader(test_data, sampler=test_sampler, batch_size=batch_size)
+    
+        model = pickle.load(open('../svm.pkl', "rb"))
+    
+    
     elif MODEL == "ct":
-        model = CTBERT.from_pretrained("ypesk/frugal-ai-ct-bert-baseline")
+        model = CTBERT.from_pretrained("ypesk/frugal-ai-EURECOM-ct-bert-baseline")
         model = model.to(device)
-        tokenizer = AutoTokenizer.from_pretrained('digitalepidemiologylab/covid-twitter-bert')
+        tokenizer = AutoTokenizer.from_pretrained('digitalepidemiologylab/covid-twitter-bert-fullset')
         
         test_texts = [t['quote'] for t in test_dataset]
     
@@ -220,7 +225,7 @@ async def evaluate_text(request: TextEvaluationRequest):
         test_dataloader = DataLoader(test_data, sampler=test_sampler, batch_size=batch_size)
     
     elif MODEL == "modern-base":
-        model = conspiracyModelBase.from_pretrained("ypesk/frugal-ai-modern-base-baseline")
+        model = conspiracyModelBase.from_pretrained("ypesk/frugal-ai-EURECOM-modern-base-fullset")
         model = model.to(device)
         tokenizer = AutoTokenizer.from_pretrained("answerdotai/ModernBERT-base")
             
@@ -241,7 +246,7 @@ async def evaluate_text(request: TextEvaluationRequest):
         test_dataloader = DataLoader(test_data, sampler=test_sampler, batch_size=batch_size)
     
     elif MODEL == "modern-large":
-        model = conspiracyModelLarge.from_pretrained("ypesk/frugal-ai-modern-large-baseline")
+        model = conspiracyModelLarge.from_pretrained('ypesk/frugal-ai-EURECOM-modern-large-fullset')
         model = model.to(device)
         tokenizer = AutoTokenizer.from_pretrained("answerdotai/ModernBERT-large")
             
@@ -262,7 +267,7 @@ async def evaluate_text(request: TextEvaluationRequest):
         test_dataloader = DataLoader(test_data, sampler=test_sampler, batch_size=batch_size)
     
     elif MODEL == "gte-base":
-        model = gteModel.from_pretrained("ypesk/frugal-ai-gte-base-baseline")
+        model = gteModel.from_pretrained("ypesk/frugal-ai-EURECOM-gte-base-fullset")
         model = model.to(device)
         tokenizer = AutoTokenizer.from_pretrained('Alibaba-NLP/gte-base-en-v1.5')
             
@@ -284,7 +289,7 @@ async def evaluate_text(request: TextEvaluationRequest):
         test_dataloader = DataLoader(test_data, sampler=test_sampler, batch_size=batch_size)
     
     elif MODEL == "gte-large":
-        model = gteModel.from_pretrained("ypesk/frugal-ai-gte-large-baseline")
+        model = gteModelLarge.from_pretrained("ypesk/frugal-ai-EURECOM-gte-large-fullset")
         model = model.to(device)
         tokenizer = AutoTokenizer.from_pretrained('Alibaba-NLP/gte-large-en-v1.5')
             
@@ -333,8 +338,7 @@ async def evaluate_text(request: TextEvaluationRequest):
                 logits = model(b_input_ids, b_input_mask, b_token_type_ids)
                 
         logits = logits.detach().cpu().numpy()
-        predictions.extend(logits.argmax(1))
-        
+        predictions.extend(logits.argmax(1))        
     
     true_labels = test_dataset["label"]   
     # Make random predictions (placeholder for actual model inference)