Upload 17 files

50731c6 verified 7 months ago

5.19 kB

#!/usr/bin/env python3
"""Export Cohere Transcribe encoder (with projection) to CoreML.
This exports the Conformer encoder + encoder_decoder_proj layer as a single model.
"""

import argparse
import sys
from pathlib import Path

import coremltools as ct
import numpy as np
import torch
import torch.nn as nn
from transformers import AutoModelForSpeechSeq2Seq


class EncoderWrapper(nn.Module):
    """Wrapper that combines encoder + projection layer."""

    def __init__(self, encoder, encoder_decoder_proj):
        super().__init__()
        self.encoder = encoder
        self.encoder_decoder_proj = encoder_decoder_proj

    def forward(self, input_features, feature_length):
        """
        Args:
            input_features: (batch, n_mels, n_frames) mel spectrogram
            feature_length: (batch,) int32 - actual length before padding
        Returns:
            hidden_states: (batch, encoded_frames, decoder_hidden_size) - encoder output after projection
        """
        encoder_outputs = self.encoder(
            input_features=input_features,
            lengths=feature_length,
            return_dict=True
        )

        hidden_states = encoder_outputs.last_hidden_state

        # Apply projection if it exists
        if self.encoder_decoder_proj is not None:
            hidden_states = self.encoder_decoder_proj(hidden_states)

        return hidden_states


def export_encoder(output_dir: Path, precision: str = "float16"):
    """Export the Cohere encoder to CoreML."""
    print("="*70)
    print("Cohere Transcribe Encoder Export")
    print("="*70)

    # Create output directory
    output_dir.mkdir(parents=True, exist_ok=True)

    # Load full model
    print("\n[1/5] Loading model from HuggingFace...")
    model = AutoModelForSpeechSeq2Seq.from_pretrained(
        "CohereLabs/cohere-transcribe-03-2026",
        trust_remote_code=True,
        torch_dtype=torch.float32,
    )
    model.eval()
    print("   ✓ Model loaded")

    # Wrap encoder + projection
    print("\n[2/5] Wrapping encoder...")
    wrapped_encoder = EncoderWrapper(model.encoder, model.encoder_decoder_proj)
    wrapped_encoder.eval()
    print("   ✓ Encoder wrapped")

    # Create example inputs
    print("\n[3/5] Creating example inputs...")
    batch_size = 1
    n_mels = 128
    max_frames = 3001  # From manifest

    example_input_features = torch.randn(batch_size, n_mels, max_frames)
    example_feature_length = torch.tensor([max_frames], dtype=torch.int32)

    print(f"   Input features: {example_input_features.shape}")
    print(f"   Feature length: {example_feature_length.shape}")

    # Trace the model
    print("\n[4/5] Tracing encoder...")
    with torch.no_grad():
        traced_encoder = torch.jit.trace(
            wrapped_encoder,
            (example_input_features, example_feature_length),
            check_trace=False,  # Disable due to conditional logic
        )

    # Test traced model
    output = traced_encoder(example_input_features, example_feature_length)
    print(f"   Output shape: {output.shape}")

    # Convert to CoreML
    print(f"\n[5/5] Converting to CoreML ({precision})...")

    # Define inputs
    inputs = [
        ct.TensorType(name="input_features", shape=example_input_features.shape, dtype=np.float32),
        ct.TensorType(name="feature_length", shape=example_feature_length.shape, dtype=np.int32),
    ]

    # Set compute precision
    compute_precision = ct.precision.FLOAT16 if precision == "float16" else ct.precision.FLOAT32

    # Convert
    mlmodel = ct.convert(
        traced_encoder,
        inputs=inputs,
        outputs=[ct.TensorType(name="hidden_states")],
        minimum_deployment_target=ct.target.iOS17,
        compute_precision=compute_precision,
    )

    # Save
    output_path = output_dir / "cohere_encoder.mlpackage"
    mlmodel.save(str(output_path))

    print(f"   ✓ Saved to: {output_path}")
    print(f"   Model size: {sum(f.stat().st_size for f in output_path.rglob('*') if f.is_file()) / 1024**3:.2f} GB")

    print("\n" + "="*70)
    print("ENCODER EXPORT COMPLETE")
    print("="*70)
    print(f"\nOutput: {output_path}")
    print(f"\nModel inputs:")
    print(f"  - input_features: (1, 128, 3001) float32 - mel spectrogram")
    print(f"  - feature_length: (1,) int32 - actual length before padding")
    print(f"\nModel output:")
    print(f"  - hidden_states: (1, 376, 1024) float16/32 - encoder output after projection")
    print()


def main():
    parser = argparse.ArgumentParser(description="Export Cohere encoder to CoreML")
    parser.add_argument(
        "--output-dir",
        type=Path,
        default=Path("build"),
        help="Output directory for CoreML models"
    )
    parser.add_argument(
        "--precision",
        choices=["float16", "float32"],
        default="float16",
        help="Model precision (default: float16)"
    )

    args = parser.parse_args()

    try:
        export_encoder(args.output_dir, args.precision)
    except Exception as e:
        print(f"\n❌ Export failed: {e}", file=sys.stderr)
        import traceback
        traceback.print_exc()
        sys.exit(1)


if __name__ == "__main__":
    main()

	#!/usr/bin/env python3
	"""Export Cohere Transcribe encoder (with projection) to CoreML.

	This exports the Conformer encoder + encoder_decoder_proj layer as a single model.
	"""

	import argparse
	import sys
	from pathlib import Path

	import coremltools as ct
	import numpy as np
	import torch
	import torch.nn as nn
	from transformers import AutoModelForSpeechSeq2Seq


	class EncoderWrapper(nn.Module):
	"""Wrapper that combines encoder + projection layer."""

	def __init__(self, encoder, encoder_decoder_proj):
	super().__init__()
	self.encoder = encoder
	self.encoder_decoder_proj = encoder_decoder_proj

	def forward(self, input_features, feature_length):
	"""
	Args:
	input_features: (batch, n_mels, n_frames) mel spectrogram
	feature_length: (batch,) int32 - actual length before padding

	Returns:
	hidden_states: (batch, encoded_frames, decoder_hidden_size) - encoder output after projection
	"""
	encoder_outputs = self.encoder(
	input_features=input_features,
	lengths=feature_length,
	return_dict=True
	)

	hidden_states = encoder_outputs.last_hidden_state

	# Apply projection if it exists
	if self.encoder_decoder_proj is not None:
	hidden_states = self.encoder_decoder_proj(hidden_states)

	return hidden_states


	def export_encoder(output_dir: Path, precision: str = "float16"):
	"""Export the Cohere encoder to CoreML."""
	print("="*70)
	print("Cohere Transcribe Encoder Export")
	print("="*70)

	# Create output directory
	output_dir.mkdir(parents=True, exist_ok=True)

	# Load full model
	print("\n[1/5] Loading model from HuggingFace...")
	model = AutoModelForSpeechSeq2Seq.from_pretrained(
	"CohereLabs/cohere-transcribe-03-2026",
	trust_remote_code=True,
	torch_dtype=torch.float32,
	)
	model.eval()
	print(" ✓ Model loaded")

	# Wrap encoder + projection
	print("\n[2/5] Wrapping encoder...")
	wrapped_encoder = EncoderWrapper(model.encoder, model.encoder_decoder_proj)
	wrapped_encoder.eval()
	print(" ✓ Encoder wrapped")

	# Create example inputs
	print("\n[3/5] Creating example inputs...")
	batch_size = 1
	n_mels = 128
	max_frames = 3001 # From manifest

	example_input_features = torch.randn(batch_size, n_mels, max_frames)
	example_feature_length = torch.tensor([max_frames], dtype=torch.int32)

	print(f" Input features: {example_input_features.shape}")
	print(f" Feature length: {example_feature_length.shape}")

	# Trace the model
	print("\n[4/5] Tracing encoder...")
	with torch.no_grad():
	traced_encoder = torch.jit.trace(
	wrapped_encoder,
	(example_input_features, example_feature_length),
	check_trace=False, # Disable due to conditional logic
	)

	# Test traced model
	output = traced_encoder(example_input_features, example_feature_length)
	print(f" Output shape: {output.shape}")

	# Convert to CoreML
	print(f"\n[5/5] Converting to CoreML ({precision})...")

	# Define inputs
	inputs = [
	ct.TensorType(name="input_features", shape=example_input_features.shape, dtype=np.float32),
	ct.TensorType(name="feature_length", shape=example_feature_length.shape, dtype=np.int32),
	]

	# Set compute precision
	compute_precision = ct.precision.FLOAT16 if precision == "float16" else ct.precision.FLOAT32

	# Convert
	mlmodel = ct.convert(
	traced_encoder,
	inputs=inputs,
	outputs=[ct.TensorType(name="hidden_states")],
	minimum_deployment_target=ct.target.iOS17,
	compute_precision=compute_precision,
	)

	# Save
	output_path = output_dir / "cohere_encoder.mlpackage"
	mlmodel.save(str(output_path))

	print(f" ✓ Saved to: {output_path}")
	print(f" Model size: {sum(f.stat().st_size for f in output_path.rglob('') if f.is_file()) / 1024*3:.2f} GB")

	print("\n" + "="*70)
	print("ENCODER EXPORT COMPLETE")
	print("="*70)
	print(f"\nOutput: {output_path}")
	print(f"\nModel inputs:")
	print(f" - input_features: (1, 128, 3001) float32 - mel spectrogram")
	print(f" - feature_length: (1,) int32 - actual length before padding")
	print(f"\nModel output:")
	print(f" - hidden_states: (1, 376, 1024) float16/32 - encoder output after projection")
	print()


	def main():
	parser = argparse.ArgumentParser(description="Export Cohere encoder to CoreML")
	parser.add_argument(
	"--output-dir",
	type=Path,
	default=Path("build"),
	help="Output directory for CoreML models"
	)
	parser.add_argument(
	"--precision",
	choices=["float16", "float32"],
	default="float16",
	help="Model precision (default: float16)"
	)

	args = parser.parse_args()

	try:
	export_encoder(args.output_dir, args.precision)
	except Exception as e:
	print(f"\n❌ Export failed: {e}", file=sys.stderr)
	import traceback
	traceback.print_exc()
	sys.exit(1)


	if __name__ == "__main__":
	main()