feat: distilbert model

requirements.txt

@@ -7,4 +7,8 @@ pydantic>=1.10.0
 python-dotenv>=1.0.0
 gradio>=4.0.0
 requests>=2.31.0
-librosa==0.10.2.post1
+librosa==0.10.2.post1
+llvmlite == 0.43.0
+transformers
+torch
+transformers[torch]

tasks/text.py

@@ -59,7 +59,51 @@ async def evaluate_text(request: TextEvaluationRequest):
 
     # Make random predictions (placeholder for actual model inference)
     true_labels = test_dataset["label"]
-    predictions = [random.randint(0, 7) for _ in range(len(true_labels))]
+
+    import torch
+    from transformers import (
+        AutoModelForSequenceClassification,
+        AutoTokenizer,
+        Trainer,
+        TrainingArguments,
+    )
+
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    model = AutoModelForSequenceClassification.from_pretrained(model_name, num_labels=len(LABEL_MAPPING))
+
+    # Tokenize the datasets
+    def tokenize_function(examples):
+        return tokenizer(examples["quote"], padding="max_length", truncation=True)
+
+    tokenized_test_dataset = test_dataset.map(tokenize_function, batched=True)
+
+    # Set training arguments
+    training_args = TrainingArguments(
+        output_dir="./bert_classification_results",
+        evaluation_strategy="epoch",
+        save_strategy="epoch",
+        learning_rate=2e-5,
+        per_device_train_batch_size=8,
+        per_device_eval_batch_size=8,
+        num_train_epochs=30,
+        weight_decay=0.01,
+        load_best_model_at_end=True,
+        use_mps_device=True,
+    )
+
+    # Initialize the Trainer
+    trainer = Trainer(
+        model=model,
+        args=training_args,
+        eval_dataset=tokenized_test_dataset,
+        tokenizer=tokenizer,
+    )
+
+    import numpy as np
+
+    preds = trainer.predict(tokenized_test_dataset)
+
+    predictions = np.array([np.argmax(x) for x in preds[0]])
 
     # --------------------------------------------------------------------------------------------
     # YOUR MODEL INFERENCE STOPS HERE