training distillbert on data (#3)

- training distillbert on data (a1e89098c1a4457b75004ffc0a21c98b10b929a5)

tasks/text.py

@@ -8,7 +8,7 @@ from .utils.evaluation import TextEvaluationRequest
 from .utils.emissions import tracker, clean_emissions_data, get_space_info
 
 #additional imports 
-from transformers import pipeline
+from transformers import Trainer, TrainingArguments, DistilBertForSequenceClassification, DistilBertTokenizerFast
 import logging
 
 router = APIRouter()
@@ -63,11 +63,37 @@ async def evaluate_text(request: TextEvaluationRequest):
     # Make random predictions (placeholder for actual model inference)
     true_labels = test_dataset["label"]
 
-    available_pipeline = pipeline(tasks="text_classfication")
-    print(available_pipeline)
-    logging.log(INFO, available_pipeline)
+    tokenizer = DistilBertTokenizerFast.from_pretrained("distilbert-base-uncased")
+    # Tokenize the datasets
+    def tokenize_function(examples):
+        return tokenizer(examples["quote"], padding="max_length", truncation=True)
+
+    train_dataset = dataset["train"].map(tokenize_function, batched=True)
+    test_dataset = dataset["test"].map(tokenize_function, batched=True)
+
+    model = DistilBertForSequenceClassification.from_pretrained("distilbert-base-uncased", num_labels=8)  # Set num_labels for your classification task
+
+    training_args = TrainingArguments(
+    output_dir="./results",            
+    eval_strategy="epoch",       # Evaluation strategy (can be "steps" or "epoch")
+    per_device_train_batch_size=16,    # Batch size for training
+    per_device_eval_batch_size=64,     # Batch size for evaluation
+    num_train_epochs=3,                # Number of training epochs
+    logging_dir="./logs",              # Directory for logs
+    logging_steps=10,                  # How often to log
+    )
     
-    predictions = [random.randint(0, 7) for _ in range(len(true_labels))]
+    trainer = Trainer(
+    model=model,                       # The model to train
+    args=training_args,                # The training arguments
+    train_dataset=train_dataset,       # The training dataset
+    eval_dataset=test_dataset          # The evaluation dataset
+    )
+    
+    
+    trainer.train()
+    predictions = trainer.evaluate()
+
 
     #--------------------------------------------------------------------------------------------
     # YOUR MODEL INFERENCE STOPS HERE