API routes. black formatting.

Finetune BERT.ipynb

@@ -62,8 +62,9 @@
     "    time_str = datetime.now().strftime(\"%Y-%m-%d %H:%M:%S\")\n",
     "    print(time_str, x)\n",
     "\n",
+    "\n",
     "class BertClassifier(nn.Module, PyTorchModelHubMixin):\n",
-    "    def __init__(self, num_labels=8, bert_variety='bert-base-uncased'):\n",
+    "    def __init__(self, num_labels=8, bert_variety=\"bert-base-uncased\"):\n",
     "        super().__init__()\n",
     "        self.bert = BertModel.from_pretrained(bert_variety)\n",
     "        self.dropout = nn.Dropout(0.05)\n",
@@ -76,6 +77,7 @@
     "        logits = self.classifier(pooled_output)\n",
     "        return logits\n",
     "\n",
+    "\n",
     "class TextDataset(Dataset):\n",
     "    def __init__(self, texts, labels, tokenizer, max_length=512):\n",
     "        self.encodings = tokenizer(\n",
@@ -83,32 +85,33 @@
     "            truncation=True,\n",
     "            padding=True,\n",
     "            max_length=max_length,\n",
-    "            return_tensors='pt',\n",
+    "            return_tensors=\"pt\",\n",
     "        )\n",
     "        self.labels = torch.tensor([int(l[0]) for l in labels])\n",
     "\n",
     "    def __getitem__(self, idx):\n",
     "        item = {key: val[idx] for key, val in self.encodings.items()}\n",
-    "        item['labels'] = self.labels[idx]\n",
+    "        item[\"labels\"] = self.labels[idx]\n",
     "        return item\n",
     "\n",
     "    def __len__(self) -> int:\n",
     "        return len(self.labels)\n",
     "\n",
+    "\n",
     "def train_model(model, train_dataloader, device, num_epochs):\n",
     "    optimizer = torch.optim.AdamW(model.parameters(), lr=2e-5)\n",
     "    criterion = nn.CrossEntropyLoss()\n",
     "    model.train()\n",
     "\n",
-    "    my_print('Starting epoch 1.')\n",
+    "    my_print(\"Starting epoch 1.\")\n",
     "    for epoch in range(num_epochs):\n",
     "        total_loss = 0\n",
     "        for batch in train_dataloader:\n",
     "            optimizer.zero_grad()\n",
     "\n",
-    "            input_ids = batch['input_ids'].to(device)\n",
-    "            attention_mask = batch['attention_mask'].to(device)\n",
-    "            labels = batch['labels'].to(device)\n",
+    "            input_ids = batch[\"input_ids\"].to(device)\n",
+    "            attention_mask = batch[\"attention_mask\"].to(device)\n",
+    "            labels = batch[\"labels\"].to(device)\n",
     "\n",
     "            outputs = model(input_ids, attention_mask)\n",
     "            loss = criterion(outputs, labels)\n",
@@ -118,7 +121,7 @@
     "\n",
     "            total_loss += loss.item()\n",
     "        avg_loss = total_loss / len(train_dataloader)\n",
-    "        my_print(f'Epoch {epoch+1}/{num_epochs} done, Average Loss: {avg_loss:0.4f}')"
+    "        my_print(f\"Epoch {epoch+1}/{num_epochs} done, Average Loss: {avg_loss:0.4f}\")"
    ]
   },
   {
@@ -137,12 +140,12 @@
    "outputs": [],
    "source": [
     "if torch.backends.mps.is_available():\n",
-    "    device = torch.device('mps')\n",
+    "    device = torch.device(\"mps\")\n",
     "    torch.mps.empty_cache()\n",
     "elif torch.cuda.is_available():\n",
-    "    device = torch.device('cuda')\n",
+    "    device = torch.device(\"cuda\")\n",
     "else:\n",
-    "    device = torch.device('cpu')"
+    "    device = torch.device(\"cpu\")"
    ]
   },
   {
@@ -162,36 +165,36 @@
    "source": [
     "def run_training(\n",
     "    max_dataset_size=16 * 200,\n",
-    "    bert_variety='bert-base-uncased',\n",
+    "    bert_variety=\"bert-base-uncased\",\n",
     "    max_length=200,\n",
     "    num_epochs=3,\n",
     "    batch_size=32,\n",
     "):\n",
     "    hf_dataset = load_dataset(\"quotaclimat/frugalaichallenge-text-train\")\n",
-    "    if not max_dataset_size == 'full' and max_dataset_size < len(hf_dataset['train']):\n",
-    "        train_dataset = hf_dataset['train'][:max_dataset_size]\n",
+    "    if not max_dataset_size == \"full\" and max_dataset_size < len(hf_dataset[\"train\"]):\n",
+    "        train_dataset = hf_dataset[\"train\"][:max_dataset_size]\n",
     "    else:\n",
-    "        train_dataset = hf_dataset['train']\n",
-    "    \n",
+    "        train_dataset = hf_dataset[\"train\"]\n",
+    "\n",
     "    tokenizer = BertTokenizer.from_pretrained(bert_variety, max_length=max_length)\n",
     "    model = BertClassifier(bert_variety=bert_variety)\n",
     "    if torch.backends.mps.is_available():\n",
-    "        device = torch.device('mps')\n",
+    "        device = torch.device(\"mps\")\n",
     "        torch.mps.empty_cache()\n",
     "    elif torch.cuda.is_available():\n",
-    "        device = torch.device('cuda')\n",
+    "        device = torch.device(\"cuda\")\n",
     "    else:\n",
-    "        device = torch.device('cpu')\n",
+    "        device = torch.device(\"cpu\")\n",
     "    model.to(device)\n",
-    "    \n",
+    "\n",
     "    dataset = TextDataset(\n",
-    "        train_dataset['quote'],\n",
-    "        train_dataset['label'],\n",
+    "        train_dataset[\"quote\"],\n",
+    "        train_dataset[\"label\"],\n",
     "        tokenizer=tokenizer,\n",
     "        max_length=max_length,\n",
     "    )\n",
     "    dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)\n",
-    "    \n",
+    "\n",
     "    train_model(model, dataloader, device, num_epochs=num_epochs)\n",
     "    return model, tokenizer"
    ]
@@ -224,7 +227,7 @@
    "source": [
     "model, tokenizer = run_training(\n",
     "    max_dataset_size=16 * 100,\n",
-    "    bert_variety='bert-base-uncased',\n",
+    "    bert_variety=\"bert-base-uncased\",\n",
     "    max_length=128,\n",
     "    num_epochs=3,\n",
     "    batch_size=32,\n",
@@ -256,27 +259,27 @@
    "source": [
     "model.eval()\n",
     "test_text = [\n",
-    "    'This was a great experience!', # 0_not_relevant\n",
-    "    'My favorite hike is Laguna de los Tres.', # 0_not_relevant\n",
-    "    'Crops will grow great in Finland if it\\'s warmer there.', # 3_not_bad\n",
-    "    'Climate change is fake.', # 1_not_happening\n",
-    "    'The apparent warming is caused by solar cycles.', # 2_not_human\n",
-    "    'Solar panels emit bad vibes.', # 4_solutions_harmful_unnecessary\n",
-    "    'All those so-called scientists are Democrats.', # 6_proponents_biased\n",
+    "    \"This was a great experience!\",  # 0_not_relevant\n",
+    "    \"My favorite hike is Laguna de los Tres.\",  # 0_not_relevant\n",
+    "    \"Crops will grow great in Finland if it's warmer there.\",  # 3_not_bad\n",
+    "    \"Climate change is fake.\",  # 1_not_happening\n",
+    "    \"The apparent warming is caused by solar cycles.\",  # 2_not_human\n",
+    "    \"Solar panels emit bad vibes.\",  # 4_solutions_harmful_unnecessary\n",
+    "    \"All those so-called scientists are Democrats.\",  # 6_proponents_biased\n",
     "]\n",
     "test_encoding = tokenizer(\n",
     "    test_text,\n",
     "    truncation=True,\n",
     "    padding=True,\n",
-    "    return_tensors='pt',\n",
+    "    return_tensors=\"pt\",\n",
     ")\n",
     "\n",
     "with torch.no_grad():\n",
-    "    test_input_ids = test_encoding['input_ids'].to(device)\n",
-    "    test_attention_mask = test_encoding['attention_mask'].to(device)\n",
+    "    test_input_ids = test_encoding[\"input_ids\"].to(device)\n",
+    "    test_attention_mask = test_encoding[\"attention_mask\"].to(device)\n",
     "    outputs = model(test_input_ids, test_attention_mask)\n",
     "    predictions = torch.argmax(outputs, dim=1)\n",
-    "    my_print(f'Predictions: {predictions}')"
+    "    my_print(f\"Predictions: {predictions}\")"
    ]
   },
   {
@@ -306,8 +309,8 @@
    ],
    "source": [
     "model, tokenizer = run_training(\n",
-    "    max_dataset_size='full',\n",
-    "    bert_variety='bert-base-uncased',\n",
+    "    max_dataset_size=\"full\",\n",
+    "    bert_variety=\"bert-base-uncased\",\n",
     "    max_length=64,\n",
     "    num_epochs=3,\n",
     "    batch_size=32,\n",
@@ -341,8 +344,8 @@
    ],
    "source": [
     "model, tokenizer = run_training(\n",
-    "    max_dataset_size='full',\n",
-    "    bert_variety='bert-base-uncased',\n",
+    "    max_dataset_size=\"full\",\n",
+    "    bert_variety=\"bert-base-uncased\",\n",
     "    max_length=128,\n",
     "    num_epochs=3,\n",
     "    batch_size=32,\n",
@@ -376,8 +379,8 @@
    ],
    "source": [
     "model, tokenizer = run_training(\n",
-    "    max_dataset_size='full',\n",
-    "    bert_variety='bert-base-uncased',\n",
+    "    max_dataset_size=\"full\",\n",
+    "    bert_variety=\"bert-base-uncased\",\n",
     "    max_length=128,\n",
     "    num_epochs=3,\n",
     "    batch_size=16,\n",
@@ -411,8 +414,8 @@
    ],
    "source": [
     "model, tokenizer = run_training(\n",
-    "    max_dataset_size='full',\n",
-    "    bert_variety='bert-base-uncased',\n",
+    "    max_dataset_size=\"full\",\n",
+    "    bert_variety=\"bert-base-uncased\",\n",
     "    max_length=256,\n",
     "    num_epochs=3,\n",
     "    batch_size=16,\n",
@@ -454,8 +457,8 @@
    ],
    "source": [
     "model_final, tokenizer_final = run_training(\n",
-    "    max_dataset_size='full',\n",
-    "    bert_variety='bert-base-uncased',\n",
+    "    max_dataset_size=\"full\",\n",
+    "    bert_variety=\"bert-base-uncased\",\n",
     "    max_length=128,\n",
     "    num_epochs=3,\n",
     "    batch_size=16,\n",
@@ -487,27 +490,27 @@
    "source": [
     "model_final.eval()\n",
     "test_text = [\n",
-    "    'This was a great experience!', # 0_not_relevant\n",
-    "    'My favorite hike is Laguna de los Tres.', # 0_not_relevant\n",
-    "    'Crops will grow great in Finland if it\\'s warmer there.', # 3_not_bad\n",
-    "    'Climate change is fake.', # 1_not_happening\n",
-    "    'The apparent warming is caused by solar cycles.', # 2_not_human\n",
-    "    'Solar panels emit bad vibes.', # 4_solutions_harmful_unnecessary\n",
-    "    'All those so-called scientists are Democrats.', # 6_proponents_biased\n",
+    "    \"This was a great experience!\",  # 0_not_relevant\n",
+    "    \"My favorite hike is Laguna de los Tres.\",  # 0_not_relevant\n",
+    "    \"Crops will grow great in Finland if it's warmer there.\",  # 3_not_bad\n",
+    "    \"Climate change is fake.\",  # 1_not_happening\n",
+    "    \"The apparent warming is caused by solar cycles.\",  # 2_not_human\n",
+    "    \"Solar panels emit bad vibes.\",  # 4_solutions_harmful_unnecessary\n",
+    "    \"All those so-called scientists are Democrats.\",  # 6_proponents_biased\n",
     "]\n",
     "test_encoding = tokenizer_final(\n",
     "    test_text,\n",
     "    truncation=True,\n",
     "    padding=True,\n",
-    "    return_tensors='pt',\n",
+    "    return_tensors=\"pt\",\n",
     ")\n",
     "\n",
     "with torch.no_grad():\n",
-    "    test_input_ids = test_encoding['input_ids'].to(device)\n",
-    "    test_attention_mask = test_encoding['attention_mask'].to(device)\n",
+    "    test_input_ids = test_encoding[\"input_ids\"].to(device)\n",
+    "    test_attention_mask = test_encoding[\"attention_mask\"].to(device)\n",
     "    outputs = model_final(test_input_ids, test_attention_mask)\n",
     "    predictions = torch.argmax(outputs, dim=1)\n",
-    "    my_print(f'Predictions: {predictions}')"
+    "    my_print(f\"Predictions: {predictions}\")"
    ]
   },
   {
@@ -550,7 +553,7 @@
     }
    ],
    "source": [
-    "model_final.push_to_hub('frugal-ai-text-bert-base')"
+    "model_final.push_to_hub(\"frugal-ai-text-bert-base\")"
    ]
   },
   {
@@ -593,7 +596,7 @@
     }
    ],
    "source": [
-    "tokenizer_final.push_to_hub('frugal-ai-text-bert-base')"
+    "tokenizer_final.push_to_hub(\"frugal-ai-text-bert-base\")"
    ]
   },
   {

app.py

@@ -7,7 +7,7 @@ load_dotenv()
 
 app = FastAPI(
     title="Frugal AI Challenge API",
-    description="API for the Frugal AI Challenge evaluation endpoints"
+    description="API for the Frugal AI Challenge evaluation endpoints",
 )
 
 # Include all routers
@@ -15,6 +15,7 @@ app.include_router(text.router)
 app.include_router(image.router)
 app.include_router(audio.router)
 
+
 @app.get("/")
 async def root():
     return {
@@ -22,6 +23,6 @@ async def root():
         "endpoints": {
             "text": "/text - Text classification task",
             "image": "/image - Image classification task (coming soon)",
-            "audio": "/audio - Audio classification task (coming soon)"
-        }
-    } 
+            "audio": "/audio - Audio classification task (coming soon)",
+        },
+    }

tasks/audio.py

@@ -9,6 +9,7 @@ from .utils.evaluation import AudioEvaluationRequest
 from .utils.emissions import tracker, clean_emissions_data, get_space_info
 
 from dotenv import load_dotenv
+
 load_dotenv()
 
 router = APIRouter()
@@ -17,13 +18,11 @@ DESCRIPTION = "Random Baseline"
 ROUTE = "/audio"
 
 
-
-@router.post(ROUTE, tags=["Audio Task"],
-             description=DESCRIPTION)
+@router.post(ROUTE, tags=["Audio Task"], description=DESCRIPTION)
 async def evaluate_audio(request: AudioEvaluationRequest):
     """
     Evaluate audio classification for rainforest sound detection.
-    
+
     Current Model: Random Baseline
     - Makes random predictions from the label space (0-1)
     - Used as a baseline for comparison
@@ -32,41 +31,40 @@ async def evaluate_audio(request: AudioEvaluationRequest):
     username, space_url = get_space_info()
 
     # Define the label mapping
-    LABEL_MAPPING = {
-        "chainsaw": 0,
-        "environment": 1
-    }
+    LABEL_MAPPING = {"chainsaw": 0, "environment": 1}
     # Load and prepare the dataset
     # Because the dataset is gated, we need to use the HF_TOKEN environment variable to authenticate
-    dataset = load_dataset(request.dataset_name,token=os.getenv("HF_TOKEN"))
-    
+    dataset = load_dataset(request.dataset_name, token=os.getenv("HF_TOKEN"))
+
     # Split dataset
-    train_test = dataset["train"].train_test_split(test_size=request.test_size, seed=request.test_seed)
+    train_test = dataset["train"].train_test_split(
+        test_size=request.test_size, seed=request.test_seed
+    )
     test_dataset = train_test["test"]
-    
+
     # Start tracking emissions
     tracker.start()
     tracker.start_task("inference")
-    
-    #--------------------------------------------------------------------------------------------
+
+    # --------------------------------------------------------------------------------------------
     # 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.
-    #--------------------------------------------------------------------------------------------   
-    
+    # --------------------------------------------------------------------------------------------
+
     # Make random predictions (placeholder for actual model inference)
     true_labels = test_dataset["label"]
     predictions = [random.randint(0, 1) for _ in range(len(true_labels))]
-    
-    #--------------------------------------------------------------------------------------------
+
+    # --------------------------------------------------------------------------------------------
     # YOUR MODEL INFERENCE STOPS HERE
-    #--------------------------------------------------------------------------------------------   
-    
+    # --------------------------------------------------------------------------------------------
+
     # Stop tracking emissions
     emissions_data = tracker.stop_task()
-    
+
     # Calculate accuracy
     accuracy = accuracy_score(true_labels, predictions)
-    
+
     # Prepare results dictionary
     results = {
         "username": username,
@@ -81,8 +79,8 @@ async def evaluate_audio(request: AudioEvaluationRequest):
         "dataset_config": {
             "dataset_name": request.dataset_name,
             "test_size": request.test_size,
-            "test_seed": request.test_seed
-        }
+            "test_seed": request.test_seed,
+        },
     }
-    
-    return results 
+
+    return results

tasks/image.py

@@ -10,6 +10,7 @@ from .utils.evaluation import ImageEvaluationRequest
 from .utils.emissions import tracker, clean_emissions_data, get_space_info
 
 from dotenv import load_dotenv
+
 load_dotenv()
 
 router = APIRouter()
@@ -17,6 +18,7 @@ router = APIRouter()
 DESCRIPTION = "Random Baseline"
 ROUTE = "/image"
 
+
 def parse_boxes(annotation_string):
     """Parse multiple boxes from a single annotation string.
     Each box has 5 values: class_id, x_center, y_center, width, height"""
@@ -26,39 +28,46 @@ def parse_boxes(annotation_string):
     for i in range(0, len(values), 5):
         if i + 5 <= len(values):
             # Skip class_id (first value) and take the next 4 values
-            box = values[i+1:i+5]
+            box = values[i + 1 : i + 5]
             boxes.append(box)
     return boxes
 
+
 def compute_iou(box1, box2):
     """Compute Intersection over Union (IoU) between two YOLO format boxes."""
+
     # Convert YOLO format (x_center, y_center, width, height) to corners
     def yolo_to_corners(box):
         x_center, y_center, width, height = box
-        x1 = x_center - width/2
-        y1 = y_center - height/2
-        x2 = x_center + width/2
-        y2 = y_center + height/2
+        x1 = x_center - width / 2
+        y1 = y_center - height / 2
+        x2 = x_center + width / 2
+        y2 = y_center + height / 2
         return np.array([x1, y1, x2, y2])
-    
+
     box1_corners = yolo_to_corners(box1)
     box2_corners = yolo_to_corners(box2)
-    
+
     # Calculate intersection
     x1 = max(box1_corners[0], box2_corners[0])
     y1 = max(box1_corners[1], box2_corners[1])
     x2 = min(box1_corners[2], box2_corners[2])
     y2 = min(box1_corners[3], box2_corners[3])
-    
+
     intersection = max(0, x2 - x1) * max(0, y2 - y1)
-    
+
     # Calculate union
-    box1_area = (box1_corners[2] - box1_corners[0]) * (box1_corners[3] - box1_corners[1])
-    box2_area = (box2_corners[2] - box2_corners[0]) * (box2_corners[3] - box2_corners[1])
+    box1_area = (box1_corners[2] - box1_corners[0]) * (
+        box1_corners[3] - box1_corners[1]
+    )
+    box2_area = (box2_corners[2] - box2_corners[0]) * (
+        box2_corners[3] - box2_corners[1]
+    )
     union = box1_area + box2_area - intersection
-    
+
     return intersection / (union + 1e-6)
 
+
 def compute_max_iou(true_boxes, pred_box):
     """Compute maximum IoU between a predicted box and all true boxes"""
     max_iou = 0
@@ -67,89 +76,91 @@ def compute_max_iou(true_boxes, pred_box):
         max_iou = max(max_iou, iou)
     return max_iou
 
-@router.post(ROUTE, tags=["Image Task"],
-             description=DESCRIPTION)
+
+@router.post(ROUTE, tags=["Image Task"], description=DESCRIPTION)
 async def evaluate_image(request: ImageEvaluationRequest):
     """
     Evaluate image classification and object detection for forest fire smoke.
-    
+
     Current Model: Random Baseline
     - Makes random predictions for both classification and bounding boxes
     - Used as a baseline for comparison
-    
+
     Metrics:
     - Classification accuracy: Whether an image contains smoke or not
     - Object Detection accuracy: IoU (Intersection over Union) for smoke bounding boxes
     """
     # Get space info
     username, space_url = get_space_info()
-    
+
     # Load and prepare the dataset
     dataset = load_dataset(request.dataset_name, token=os.getenv("HF_TOKEN"))
-    
+
     # Split dataset
-    train_test = dataset["train"].train_test_split(test_size=request.test_size, seed=request.test_seed)
+    train_test = dataset["train"].train_test_split(
+        test_size=request.test_size, seed=request.test_seed
+    )
     test_dataset = train_test["test"]
-    
+
     # Start tracking emissions
     tracker.start()
     tracker.start_task("inference")
-    
-    #--------------------------------------------------------------------------------------------
+
+    # --------------------------------------------------------------------------------------------
     # YOUR MODEL INFERENCE CODE HERE
     # Update the code below to replace the random baseline with your model inference
-    #--------------------------------------------------------------------------------------------   
-    
+    # --------------------------------------------------------------------------------------------
+
     predictions = []
     true_labels = []
     pred_boxes = []
     true_boxes_list = []  # List of lists, each inner list contains boxes for one image
-    
+
     for example in test_dataset:
         # Parse true annotation (YOLO format: class_id x_center y_center width height)
         annotation = example.get("annotations", "").strip()
         has_smoke = len(annotation) > 0
         true_labels.append(int(has_smoke))
-        
+
         # Make random classification prediction
         pred_has_smoke = random.random() > 0.5
         predictions.append(int(pred_has_smoke))
-        
+
         # If there's a true box, parse it and make random box prediction
         if has_smoke:
             # Parse all true boxes from the annotation
             image_true_boxes = parse_boxes(annotation)
             true_boxes_list.append(image_true_boxes)
-            
+
             # For baseline, make one random box prediction per image
             # In a real model, you might want to predict multiple boxes
             random_box = [
                 random.random(),  # x_center
                 random.random(),  # y_center
                 random.random() * 0.5,  # width (max 0.5)
-                random.random() * 0.5   # height (max 0.5)
+                random.random() * 0.5,  # height (max 0.5)
             ]
             pred_boxes.append(random_box)
-    
-    #--------------------------------------------------------------------------------------------
+
+    # --------------------------------------------------------------------------------------------
     # YOUR MODEL INFERENCE STOPS HERE
-    #--------------------------------------------------------------------------------------------   
-    
+    # --------------------------------------------------------------------------------------------
+
     # Stop tracking emissions
     emissions_data = tracker.stop_task()
-    
+
     # Calculate classification accuracy
     classification_accuracy = accuracy_score(true_labels, predictions)
-    
+
     # Calculate mean IoU for object detection (only for images with smoke)
     # For each image, we compute the max IoU between the predicted box and all true boxes
     ious = []
     for true_boxes, pred_box in zip(true_boxes_list, pred_boxes):
         max_iou = compute_max_iou(true_boxes, pred_box)
         ious.append(max_iou)
-    
+
     mean_iou = float(np.mean(ious)) if ious else 0.0
-    
+
     # Prepare results dictionary
     results = {
         "username": username,
@@ -165,8 +176,8 @@ async def evaluate_image(request: ImageEvaluationRequest):
         "dataset_config": {
             "dataset_name": request.dataset_name,
             "test_size": request.test_size,
-            "test_seed": request.test_seed
-        }
+            "test_seed": request.test_seed,
+        },
     }
-    
-    return results 
+
+    return results

tasks/text.py

@@ -12,13 +12,21 @@ from .utils.emissions import tracker, clean_emissions_data, get_space_info
 
 router = APIRouter()
 
-DESCRIPTION = "bert base finetuned"
+DESCRIPTIONS = {
+    "baseline": "baseline most common class",
+    "bert-base": "bert base finetuned",
+    "bert-medium": "to be implemented",
+    "bert-small": "to be implemented",
+    "bert-mini": "to be implemented",
+    "bert-tiny": "to be implemented",
+}
+
 ROUTE = "/text"
 
 
 def baseline_model(dataset_length: int):
     # Make random predictions (placeholder for actual model inference)
-    #predictions = [random.randint(0, 7) for _ in range(dataset_length)]
+    # predictions = [random.randint(0, 7) for _ in range(dataset_length)]
 
     # My favorate baseline is the most common class.
     predictions = [0] * dataset_length
@@ -26,48 +34,50 @@ def baseline_model(dataset_length: int):
     return predictions
 
 
-def bert_model(test_dataset):
-    print('Starting my code block.')
+def bert_model(test_dataset: dict, model_type: str):
+    print("Starting my code block.")
     texts = test_dataset["quote"]
 
-    model_repo = 'Nonnormalizable/frugal-ai-text-bert-base'
+    model_repo = f"Nonnormalizable/frugal-ai-text-{model_type}"
     config = AutoConfig.from_pretrained(model_repo)
     model = AutoModelForSequenceClassification.from_pretrained(model_repo)
     tokenizer = AutoTokenizer.from_pretrained(model_repo)
 
     if torch.cuda.is_available():
-        device = torch.device('cuda')
+        device = torch.device("cuda")
     else:
-        device = torch.device('cpu')
-    print('device:', device)
+        device = torch.device("cpu")
+    print("device:", device)
     model = model.to(device)
     test_encoding = tokenizer(
         texts,
         truncation=True,
         padding=True,
-        return_tensors='pt',
-        )
+        return_tensors="pt",
+    )
 
     model.eval()
     with torch.no_grad():
-        test_input_ids = test_encoding['input_ids'].to(device)
-        test_attention_mask = test_encoding['attention_mask'].to(device)
-        print('Starting model run.')
+        test_input_ids = test_encoding["input_ids"].to(device)
+        test_attention_mask = test_encoding["attention_mask"].to(device)
+        print("Starting model run.")
         outputs = model(test_input_ids, test_attention_mask)
-        print('End of model run.')
+        print("End of model run.")
         predictions = torch.argmax(outputs.logits, dim=1)
         predictions = predictions.cpu().numpy()
-        
-    print('End of my code block.')
+
+    print("End of my code block.")
     return predictions
 
 
-@router.post(ROUTE, tags=["Text Task"], 
-             description=DESCRIPTION)
-async def evaluate_text(request: TextEvaluationRequest):
+@router.post(ROUTE, tags=["Text Task"])
+async def evaluate_text(
+    request: TextEvaluationRequest,
+    model_type="bert-base",
+):
     """
     Evaluate text classification for climate disinformation detection.
-    
+
     Current Model: Random Baseline
     - Makes random predictions from the label space (0-7)
     - Used as a baseline for comparison
@@ -84,7 +94,7 @@ async def evaluate_text(request: TextEvaluationRequest):
         "4_solutions_harmful_unnecessary": 4,
         "5_science_unreliable": 5,
         "6_proponents_biased": 6,
-        "7_fossil_fuels_needed": 7
+        "7_fossil_fuels_needed": 7,
     }
 
     # Load and prepare the dataset
@@ -94,39 +104,44 @@ async def evaluate_text(request: TextEvaluationRequest):
     dataset = dataset.map(lambda x: {"label": LABEL_MAPPING[x["label"]]})
 
     # Split dataset
-    train_test = dataset["train"].train_test_split(test_size=request.test_size, seed=request.test_seed)
+    train_test = dataset["train"].train_test_split(
+        test_size=request.test_size, seed=request.test_seed
+    )
     test_dataset = train_test["test"]
-    
+
     # Start tracking emissions
     tracker.start()
     tracker.start_task("inference")
 
-    #--------------------------------------------------------------------------------------------
+    # --------------------------------------------------------------------------------------------
     # 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.
-    #--------------------------------------------------------------------------------------------   
+    # --------------------------------------------------------------------------------------------
 
     true_labels = test_dataset["label"]
-    #predictions = baseline_model(len(true_labels))
-    predictions = bert_model(test_dataset)
+    if model_type == "baseline":
+        predictions = baseline_model(len(true_labels))
+    elif model_type[:5] == "bert-":
+        predictions = bert_model(test_dataset, model_type)
+    else:
+        raise ValueError(model_type)
 
-    #--------------------------------------------------------------------------------------------
+    # --------------------------------------------------------------------------------------------
     # YOUR MODEL INFERENCE STOPS HERE
-    #--------------------------------------------------------------------------------------------   
+    # --------------------------------------------------------------------------------------------
 
-    
     # Stop tracking emissions
     emissions_data = tracker.stop_task()
-    
+
     # Calculate accuracy
     accuracy = accuracy_score(true_labels, predictions)
-    
+
     # Prepare results dictionary
     results = {
         "username": username,
         "space_url": space_url,
         "submission_timestamp": datetime.now().isoformat(),
-        "model_description": DESCRIPTION,
+        "model_description": DESCRIPTIONS[model_type],
         "accuracy": float(accuracy),
         "energy_consumed_wh": emissions_data.energy_consumed * 1000,
         "emissions_gco2eq": emissions_data.emissions * 1000,
@@ -135,8 +150,8 @@ async def evaluate_text(request: TextEvaluationRequest):
         "dataset_config": {
             "dataset_name": request.dataset_name,
             "test_size": request.test_size,
-            "test_seed": request.test_seed
-        }
+            "test_seed": request.test_seed,
+        },
     }
-    
-    return results 
+
+    return results

tasks/utils/emissions.py

@@ -4,17 +4,26 @@ import os
 # Initialize tracker
 tracker = EmissionsTracker(allow_multiple_runs=True)
 
+
 class EmissionsData:
     def __init__(self, energy_consumed: float, emissions: float):
         self.energy_consumed = energy_consumed
         self.emissions = emissions
 
+
 def clean_emissions_data(emissions_data):
     """Remove unwanted fields from emissions data"""
     data_dict = emissions_data.__dict__
-    fields_to_remove = ['timestamp', 'project_name', 'experiment_id', 'latitude', 'longitude']
+    fields_to_remove = [
+        "timestamp",
+        "project_name",
+        "experiment_id",
+        "latitude",
+        "longitude",
+    ]
     return {k: v for k, v in data_dict.items() if k not in fields_to_remove}
 
+
 def get_space_info():
     """Get the space username and URL from environment variables"""
     space_name = os.getenv("SPACE_ID", "")
@@ -25,4 +34,4 @@ def get_space_info():
             return username, space_url
         except Exception as e:
             print(f"Error getting space info: {e}")
-    return "local-user", "local-development" 
+    return "local-user", "local-development"

tasks/utils/evaluation.py

@@ -1,18 +1,28 @@
 from typing import Optional
 from pydantic import BaseModel, Field
 
+
 class BaseEvaluationRequest(BaseModel):
-    test_size: float = Field(0.2, ge=0.0, le=1.0, description="Size of the test split (between 0 and 1)")
+    test_size: float = Field(
+        0.2, ge=0.0, le=1.0, description="Size of the test split (between 0 and 1)"
+    )
     test_seed: int = Field(42, ge=0, description="Random seed for reproducibility")
 
+
 class TextEvaluationRequest(BaseEvaluationRequest):
-    dataset_name: str = Field("QuotaClimat/frugalaichallenge-text-train", 
-                            description="The name of the dataset on HuggingFace Hub")
+    dataset_name: str = Field(
+        "QuotaClimat/frugalaichallenge-text-train",
+        description="The name of the dataset on HuggingFace Hub",
+    )
+
 
 class ImageEvaluationRequest(BaseEvaluationRequest):
-    dataset_name: str = Field("pyronear/pyro-sdis", 
-                            description="The name of the dataset on HuggingFace Hub")
+    dataset_name: str = Field(
+        "pyronear/pyro-sdis", description="The name of the dataset on HuggingFace Hub"
+    )
+
 
 class AudioEvaluationRequest(BaseEvaluationRequest):
-    dataset_name: str = Field("rfcx/frugalai", 
-                            description="The name of the dataset on HuggingFace Hub") 
+    dataset_name: str = Field(
+        "rfcx/frugalai", description="The name of the dataset on HuggingFace Hub"
+    )