Update tasks/text.py

tasks/text.py

@@ -21,6 +21,63 @@ router = APIRouter()
 DESCRIPTION = "First Baseline"
 ROUTE = "/text"
 
+
+MODEL = "mlp" #mlp, ct, modern
+
+class ConspiracyClassification(
+    nn.Module,
+    PyTorchModelHubMixin, 
+    # optionally, you can add metadata which gets pushed to the model card
+):    
+    def __init__(self, num_classes):
+        super().__init__()
+        self.h1 = nn.Linear(384, 100)
+        self.h2 = nn.Linear(100, 100)
+        self.h3 = nn.Linear(100, 100)
+        self.h4 = nn.Linear(100, 50)
+        self.h5 = nn.Linear(50, num_classes)
+        self.dropout = nn.Dropout(0.2)
+        self.activation = nn.ReLU()
+
+        
+    def forward(self, input_texts):
+        outputs = self.h1(input_texts)
+        outputs = self.activation(outputs)
+        outputs = self.dropout(outputs)
+        outputs = self.h2(outputs)
+        outputs = self.activation(outputs)
+        outputs = self.dropout(outputs)
+        outputs = self.h3(outputs)
+        outputs = self.activation(outputs)
+        outputs = self.dropout(outputs)
+        outputs = self.h4(outputs)
+        outputs = self.activation(outputs)
+        outputs = self.dropout(outputs)
+        outputs = self.h5(outputs)
+        
+        return outputs  
+
+class CovidTwitterBertClassifier(
+    nn.Module,
+    PyTorchModelHubMixin, 
+    # optionally, you can add metadata which gets pushed to the model card
+):    
+    def __init__(self, num_classes):
+        super().__init__()
+        self.n_classes = num_classes
+        self.bert = BertForPreTraining.from_pretrained('digitalepidemiologylab/covid-twitter-bert-v2')    
+        self.bert.cls.seq_relationship = nn.Linear(1024, num_classes)
+
+        self.sigmoid = nn.Sigmoid()
+        
+    def forward(self, input_ids, token_type_ids, input_mask):
+        outputs = self.bert(input_ids = input_ids, token_type_ids = token_type_ids, attention_mask = input_mask)
+
+        logits = outputs[1]
+        
+        return logits  
+
+
 @router.post(ROUTE, tags=["Text Task"], 
              description=DESCRIPTION)
 async def evaluate_text(request: TextEvaluationRequest):
@@ -64,53 +121,41 @@ async def evaluate_text(request: TextEvaluationRequest):
     # YOUR MODEL INFERENCE CODE HERE
     # Update the code below to replace the random baseline by your model inference within the inference pass where the energy consumption and emissions are tracked.
     #--------------------------------------------------------------------------------------------   
-    class CovidTwitterBertClassifier(
-        nn.Module,
-        PyTorchModelHubMixin, 
-        # optionally, you can add metadata which gets pushed to the model card
-    ):    
-        def __init__(self, num_classes):
-            super().__init__()
-            self.n_classes = num_classes
-            self.bert = BertForPreTraining.from_pretrained('digitalepidemiologylab/covid-twitter-bert-v2')    
-            self.bert.cls.seq_relationship = nn.Linear(1024, num_classes)
-    
-            self.sigmoid = nn.Sigmoid()
-            
-        def forward(self, input_ids, token_type_ids, input_mask):
-            outputs = self.bert(input_ids = input_ids, token_type_ids = token_type_ids, attention_mask = input_mask)
-    
-            logits = outputs[1]
-            
-            return logits  
-    model = CovidTwitterBertClassifier.from_pretrained("ypesk/ct-baseline")
-    model.eval()
-
-
-    tokenizer = AutoTokenizer.from_pretrained('digitalepidemiologylab/covid-twitter-bert')
-    
-    test_texts = [t['quote'] for t in test_dataset]
+    if MODEL =="mlp":
+        model = ConspiracyClassification.from_pretrained("ypesk/frugal-ai-mlp-baseline")
+        
+        emb_model = SentenceTransformer("paraphrase-MiniLM-L3-v2")
+        batch_size = 6
 
-    MAX_LEN = 256 #1024 # < m some tweets will be truncated
-    
-    tokenized_test = tokenizer(test_texts, max_length=MAX_LEN, padding='max_length', truncation=True)
-    test_input_ids, test_token_type_ids, test_attention_mask = tokenized_test['input_ids'], tokenized_test['token_type_ids'], tokenized_test['attention_mask']
-    test_token_type_ids = torch.tensor(test_token_type_ids)
+        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 == "ct":
+        model = CovidTwitterBertClassifier.from_pretrained("ypesk/ct-baseline")
+        tokenizer = AutoTokenizer.from_pretrained('digitalepidemiologylab/covid-twitter-bert')
+        
+        test_texts = [t['quote'] for t in test_dataset]
     
-    test_input_ids = torch.tensor(test_input_ids)
-    test_attention_mask = torch.tensor(test_attention_mask)
-
-    batch_size = 12 #
-    test_data = TensorDataset(test_input_ids, test_attention_mask, test_token_type_ids)
+        MAX_LEN = 256 #1024 # < m some tweets will be truncated
+        
+        tokenized_test = tokenizer(test_texts, max_length=MAX_LEN, padding='max_length', truncation=True)
+        test_input_ids, test_token_type_ids, test_attention_mask = tokenized_test['input_ids'], tokenized_test['token_type_ids'], tokenized_test['attention_mask']
+        test_token_type_ids = torch.tensor(test_token_type_ids)
+        
+        test_input_ids = torch.tensor(test_input_ids)
+        test_attention_mask = torch.tensor(test_attention_mask)
     
-    test_sampler = SequentialSampler(test_data)
-    test_dataloader = DataLoader(test_data, sampler=test_sampler, batch_size=batch_size)
+        batch_size = 12 #
+        test_data = TensorDataset(test_input_ids, test_attention_mask, test_token_type_ids)
+        
+        test_sampler = SequentialSampler(test_data)
+        test_dataloader = DataLoader(test_data, sampler=test_sampler, batch_size=batch_size)
 
+    model.eval()
     predictions = []
-    c=0
     for batch in tqdm(test_dataloader):
-        print(c)
-        c+=1
 
         b_input_ids, b_input_mask, b_token_type_ids = batch
         with torch.no_grad():