Create train_acft.py

train_acft.py

@@ -0,0 +1,201 @@
+import argparse
+import os
+import torch
+from tqdm import tqdm
+from torch import nn
+from datasets import load_dataset, Audio
+from transformers import WhisperModel, WhisperProcessor, get_linear_schedule_with_warmup
+from torch.utils.tensorboard import SummaryWriter
+from dotenv import load_dotenv
+
+load_dotenv()
+HF_TOKEN = os.environ.get("HF_TOKEN")
+
+class SlicedEmbedding(nn.Module):
+    def __init__(self, orig_embed, n_ctx):
+        super().__init__()
+        self.orig_embed_ref = [orig_embed]
+        self.n_ctx = n_ctx
+        self.num_embeddings = n_ctx
+        
+    @property
+    def weight(self):
+        return self.orig_embed_ref[0].weight[:self.n_ctx]
+        
+    def forward(self, input_ids):
+        return self.orig_embed_ref[0](input_ids)
+
+def get_sample(example, processor):
+    waveform = example["audio"]["array"]
+    sampling_rate = example["audio"]["sampling_rate"]
+
+    input_features = processor(
+        waveform, sampling_rate=sampling_rate, return_tensors="pt"
+    ).input_features
+
+    return {
+        "length": len(waveform) / sampling_rate,
+        "input_features": input_features,
+        "input_ids": processor.tokenizer.encode(example["text"].lower())
+    }
+
+def compute_partially_encoder(model, data, n_audio_ctx):
+    diffy = 2*n_audio_ctx - data.shape[2]
+
+    if diffy > 0:
+        data = nn.functional.pad(data, [0, diffy, 0, 0, 0, 0], "constant", 0.0)
+    elif diffy < 0:
+        data = data[:,:,:diffy]
+
+    if n_audio_ctx == 1500:
+        return model.encoder(data).last_hidden_state
+
+    orig_embed = model.encoder.embed_positions
+    orig_max_pos = model.encoder.config.max_source_positions
+    
+    model.encoder.embed_positions = SlicedEmbedding(orig_embed, n_ctx=n_audio_ctx)
+    model.encoder.config.max_source_positions = n_audio_ctx
+    
+    try:
+        output = model.encoder(data).last_hidden_state
+    finally:
+        model.encoder.embed_positions = orig_embed
+        model.encoder.config.max_source_positions = orig_max_pos
+        
+    return output
+
+def compute_hidden_state_loss(model_train, model_base, criterion, example):
+    n_ctx = int(round((1500.0 / 30.0) * example["length"] ))
+    
+    assert 0 < n_ctx <= 1500, f"Invalid n_ctx calculated: {n_ctx}"
+
+    extra_ctx = torch.randint(-min(64, n_ctx // 3), min(64, n_ctx // 3), (1,)).item()
+    n_ctx += extra_ctx
+    n_ctx = max(1, min(1500, n_ctx))
+
+    input_features = example["input_features"].cuda()
+    input_ids = torch.tensor([example["input_ids"]], dtype=torch.long).cuda()
+
+    encoder_hidden_states_partial = compute_partially_encoder(model_train, input_features, n_ctx)
+    
+    output_partial = model_train.decoder(
+        input_ids=input_ids,
+        encoder_hidden_states=encoder_hidden_states_partial,
+        output_hidden_states=True
+    )
+
+    with torch.no_grad():
+        encoder_hidden_states_full = compute_partially_encoder(model_base, input_features, 1500)
+        output_full = model_base.decoder(
+            input_ids=input_ids,
+            encoder_hidden_states=encoder_hidden_states_full,
+            output_hidden_states=True
+        )
+        
+    student_tensors = torch.cat(output_partial.hidden_states, 0)
+    teacher_tensors = torch.cat(output_full.hidden_states, 0)
+
+    loss = criterion(student_tensors, teacher_tensors)
+    loss.backward()
+
+    return loss.item()
+
+def save_checkpoint(model_train, size, processor, output_dir):
+    from transformers import WhisperForConditionalGeneration
+    final_model = WhisperForConditionalGeneration.from_pretrained(f"openai/whisper-{size}").eval().cpu()
+    final_model.model = model_train.eval().cpu()
+    final_model.save_pretrained(output_dir)
+    processor.save_pretrained(output_dir)
+    model_train.cuda().train()
+
+def train_futo_script(size):
+    print(f"Starting exact FUTO distillation for model: {size}")
+    param_counts = {"tiny": "39M", "base": "74M", "small": "244M"}
+    
+    model_train = WhisperModel.from_pretrained(f"openai/whisper-{size}").cuda().train()
+    model_base = WhisperModel.from_pretrained(f"openai/whisper-{size}").cuda().eval()
+    
+    processor = WhisperProcessor.from_pretrained(f"openai/whisper-small", language="danish", task="transcribe")
+    
+    ds = load_dataset("CoRal-project/coral-v3", "read_aloud", token=HF_TOKEN, split="train", streaming=True)
+    ds = ds.cast_column("audio", Audio(sampling_rate=16000))
+    
+    criterion = torch.nn.MSELoss()
+    
+    # Hyperparameters
+    learning_rate = 1e-6
+    weight_decay = 0.1
+    max_training_steps = 20000
+
+    optimizer = torch.optim.AdamW(model_train.parameters(), lr=learning_rate, weight_decay=weight_decay)
+    
+    writer = SummaryWriter()
+    writer.add_text("name", f"{size} v3")
+    
+    num_length = 0
+    step = 0
+    running_loss = 0.0
+    
+    best_loss = float('inf')
+    patience = 20
+    patience_counter = 0
+    eval_interval = 500
+    
+    pbar = tqdm(ds)
+    try:
+        for raw_example in pbar:
+            duration = len(raw_example["audio"]["array"]) / 16000.0
+            if duration > 29.0: 
+                continue
+                
+            example = get_sample(raw_example, processor)
+            
+            optimizer.zero_grad()
+            
+            # Compute loss and immediately update (Batch Size 1)
+            loss_val = compute_hidden_state_loss(model_train, model_base, criterion, example)
+            optimizer.step()
+            
+            step += 1
+            num_length += example["length"]
+            
+            # Update EMA loss
+            running_loss = loss_val if step == 1 else 0.95 * running_loss + 0.05 * loss_val
+            
+            writer.add_scalar("loss/train", loss_val, step)
+            writer.add_scalar("length/train", num_length, step)
+            
+            pbar.set_description(f"Step {step}, Avg Loss: {running_loss:.4f}")
+
+            # Checkpoint
+            if step % eval_interval == 0:
+                if running_loss < best_loss:
+                    best_loss = running_loss
+                    patience_counter = 0
+                    checkpoint_dir = f"{size}_{param_counts.get(size, 'unknown')}_danish_whisper_acft_futo_best"
+                    save_checkpoint(model_train, size, processor, checkpoint_dir)
+                    tqdm.write(f"\n[Step {step}] New best loss: {best_loss:.4f}. Saved checkpoint to {checkpoint_dir}")
+                else:
+                    patience_counter += 1
+                    tqdm.write(f"\n[Step {step}] No improvement. Patience: {patience_counter}/{patience}")
+                    if patience_counter >= patience:
+                        tqdm.write("\n[Early Stopping] Loss hasn't improved. Halting training.")
+                        break
+
+            if step >= max_training_steps:
+                tqdm.write("\n[Max Steps Reached] Halting training.")
+                break
+
+    except KeyboardInterrupt:
+        print("\n\n[CTRL+C detected] Training manually interrupted! Proceeding to save the final model...")
+
+    output_dir = f"{size}_{param_counts.get(size, 'unknown')}_danish_whisper_acft_futo_latest"
+    print(f"\nSaving latest model to {output_dir}")
+    save_checkpoint(model_train, size, processor, output_dir)
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Run exact FUTO script structure.")
+    parser.add_argument("--size", choices=["tiny", "base", "small"], default="base")
+    args = parser.parse_args()
+
+    train_futo_script(args.size)