Update train_optillm_classifier.py

bert large for classification

Author: codelion

scripts/train_optillm_classifier.py

@@ -24,17 +24,17 @@
 MAX_LENGTH = 512
 
 class OptILMDataset(Dataset):
-    def __init__(self, prompts, approaches, tokenizer):
+    def __init__(self, prompts, best_approaches, tokenizer):
         self.prompts = prompts
-        self.approaches = approaches
+        self.best_approaches = best_approaches
         self.tokenizer = tokenizer
 
     def __len__(self):
         return len(self.prompts)
 
     def __getitem__(self, idx):
         prompt = self.prompts[idx]
-        approach = self.approaches[idx]
+        best_approach = self.best_approaches[idx]
 
         encoding = self.tokenizer.encode_plus(
             prompt,
@@ -49,69 +49,60 @@ def __getitem__(self, idx):
         return {
             'input_ids': encoding['input_ids'].flatten(),
             'attention_mask': encoding['attention_mask'].flatten(),
-            'labels': torch.tensor(APPROACHES.index(approach), dtype=torch.long)
+            'labels': torch.tensor(APPROACHES.index(best_approach), dtype=torch.long)
         }
 
 def load_and_preprocess_data(tokenizer):
-    dataset = load_dataset('json', data_files='optillm_dataset_1.jsonl')
+    dataset = load_dataset('json', data_files='optillm_dataset.jsonl')
 
     data_items = []
 
     for item in dataset['train']:
         prompt = item['prompt']
         results = item['results']
 
-        valid_results = [r for r in results if 'approach' in r]
-        if not valid_results:
+        if not results:
             continue
+        # Filter the list to exclude items where rank is None
+        filtered_data = [item for item in results if item['rank'] is not None]
+        # Find the best approach (lowest rank)
+        best_result = min(filtered_data, key=lambda x: x['rank'])
+        best_approach = best_result['approach']
 
-        for result in valid_results:
-            data_items.append({
-                'prompt': prompt,
-                'approach': result['approach']
-            })
+        data_items.append({
+            'prompt': prompt,
+            'best_approach': best_approach
+        })
 
     # Print some statistics
     print(f"Total data points: {len(data_items)}")
     print(f"Unique prompts: {len(set(item['prompt'] for item in data_items))}")
-    approach_counts = Counter(item['approach'] for item in data_items)
-    print("Approach distribution:")
+    approach_counts = Counter(item['best_approach'] for item in data_items)
+    print("Best Approach distribution:")
     for approach, count in approach_counts.items():
         print(f"  {approach}: {count}")
 
-    # Calculate class weights for balanced sampling
-    class_weights = {approach: len(data_items) / count for approach, count in approach_counts.items()}
-    sample_weights = [class_weights[item['approach']] for item in data_items]
-
     # Split the data
     train_data, val_data = train_test_split(data_items, test_size=0.2, random_state=42)
 
     train_dataset = OptILMDataset(
         [item['prompt'] for item in train_data],
-        [item['approach'] for item in train_data],
+        [item['best_approach'] for item in train_data],
         tokenizer
     )
     val_dataset = OptILMDataset(
         [item['prompt'] for item in val_data],
-        [item['approach'] for item in val_data],
+        [item['best_approach'] for item in val_data],
         tokenizer
     )
 
-    # Create a weighted sampler for the training data
-    train_sampler = WeightedRandomSampler(
-        weights=[class_weights[item['approach']] for item in train_data],
-        num_samples=len(train_data),
-        replacement=True
-    )
-
-    return train_dataset, val_dataset, train_sampler
+    return train_dataset, val_dataset
 
-def calculate_accuracy(logits, labels):
-    predictions = torch.argmax(logits, dim=-1)
+def calculate_accuracy(predictions, labels):
     return (predictions == labels).float().mean()
 
 def train(model, train_dataloader, val_dataloader, optimizer, scheduler, num_epochs):
-    best_val_loss = float('inf')
+    best_val_accuracy = 0.0
 
     for epoch in range(num_epochs):
         model.train()
@@ -134,14 +125,14 @@ def train(model, train_dataloader, val_dataloader, optimizer, scheduler, num_epo
             optimizer.zero_grad()
 
             total_loss += loss.item()
-            total_accuracy += calculate_accuracy(logits, labels)
+            predictions = torch.argmax(logits, dim=-1)
+            total_accuracy += calculate_accuracy(predictions, labels)
 
         avg_train_loss = total_loss / len(train_dataloader)
         avg_train_accuracy = total_accuracy / len(train_dataloader)
 
         # Validation
         model.eval()
-        total_val_loss = 0
         total_val_accuracy = 0
 
         with torch.no_grad():
@@ -150,20 +141,17 @@ def train(model, train_dataloader, val_dataloader, optimizer, scheduler, num_epo
                 attention_mask = batch['attention_mask'].to(device)
                 labels = batch['labels'].to(device)
 
-                outputs = model(input_ids, attention_mask=attention_mask, labels=labels)
-                val_loss = outputs.loss
+                outputs = model(input_ids, attention_mask=attention_mask)
                 logits = outputs.logits
+                predictions = torch.argmax(logits, dim=-1)
+                total_val_accuracy += calculate_accuracy(predictions, labels)
 
-                total_val_loss += val_loss.item()
-                total_val_accuracy += calculate_accuracy(logits, labels)
-
-        avg_val_loss = total_val_loss / len(val_dataloader)
         avg_val_accuracy = total_val_accuracy / len(val_dataloader)
 
-        print(f"Epoch {epoch+1}/{num_epochs}, Train Loss: {avg_train_loss:.4f}, Train Accuracy: {avg_train_accuracy:.4f}, Val Loss: {avg_val_loss:.4f}, Val Accuracy: {avg_val_accuracy:.4f}")
+        print(f"Epoch {epoch+1}/{num_epochs}, Train Loss: {avg_train_loss:.4f}, Train Accuracy: {avg_train_accuracy:.4f}, Val Accuracy: {avg_val_accuracy:.4f}")
 
-        if avg_val_loss < best_val_loss:
-            best_val_loss = avg_val_loss
+        if avg_val_accuracy > best_val_accuracy:
+            best_val_accuracy = avg_val_accuracy
             # Save the best model
             save_model(model, "best_model.safetensors")
 
@@ -188,10 +176,10 @@ def main(args):
     model.to(device)
 
     # Load and preprocess data
-    train_dataset, val_dataset, train_sampler = load_and_preprocess_data(tokenizer)
+    train_dataset, val_dataset = load_and_preprocess_data(tokenizer)
 
     # Create data loaders
-    train_dataloader = DataLoader(train_dataset, batch_size=args.batch_size, sampler=train_sampler)
+    train_dataloader = DataLoader(train_dataset, batch_size=args.batch_size)
     val_dataloader = DataLoader(val_dataset, batch_size=args.batch_size)
 
     # Optimizer and scheduler
@@ -209,17 +197,31 @@ def main(args):
         model.push_to_hub(args.hub_model_id)
         tokenizer.push_to_hub(args.hub_model_id)
 
-    # Example inference
-    test_prompt = "Maximize x + y subject to: x + 2y <= 10, x >= 0, y >= 0"
-    predicted_approach, confidence = inference(model, tokenizer, test_prompt)
-    print(f"Test Prompt: {test_prompt}")
-    print(f"Predicted Approach: {predicted_approach}")
-    print(f"Confidence: {confidence:.4f}")
+    # Example inferences
+    test_prompts = [
+        "Maximize x + y subject to: x + 2y <= 10, x >= 0, y >= 0",
+        "Find the shortest path between nodes A and B in the given graph",
+        "Solve the Tower of Hanoi problem with 4 disks",
+        "Determine if the given number is prime",
+        "Find all possible combinations of coins that sum up to $1",
+        "Implement a binary search algorithm",
+        "Design an algorithm to find the longest palindromic substring",
+        "Solve the 8-queens problem",
+        "Implement a depth-first search algorithm for a graph",
+        "Find the maximum subarray sum in a given array of integers"
+    ]
+
+    print("\nInference Examples:")
+    for prompt in test_prompts:
+        predicted_approach, confidence = inference(model, tokenizer, prompt)
+        print(f"\nTest Prompt: {prompt}")
+        print(f"Predicted Approach: {predicted_approach}")
+        print(f"Confidence: {confidence:.4f}")
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(description="Train OptILM classifier")
-    parser.add_argument("--model_name", type=str, default="roberta-large", help="Pretrained model name")
-    parser.add_argument("--batch_size", type=int, default=16, help="Batch size for training")
+    parser.add_argument("--model_name", type=str, default="google-bert/bert-large-uncased", help="Pretrained model name")
+    parser.add_argument("--batch_size", type=int, default=4, help="Batch size for training")
     parser.add_argument("--learning_rate", type=float, default=2e-5, help="Learning rate")
     parser.add_argument("--num_epochs", type=int, default=10, help="Number of training epochs")
     parser.add_argument("--push_to_hub", action="store_true", help="Push model to Hugging Face Hub")