File size: 1,248 Bytes
4e6f878 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 |
import torch
import numpy as np
from transformers import AutoModelForCausalLM, AutoTokenizer
import coremltools as ct
# Load model and convert to TorchScript
model = AutoModelForCausalLM.from_pretrained("LSX-UniWue/LLaMmlein_1B")
tokenizer = AutoTokenizer.from_pretrained("LSX-UniWue/LLaMmlein_1B")
# Set model to eval mode
model.eval()
# Create example input
text = "Ein Beispieltext"
inputs = tokenizer(text, return_tensors="pt")
# Create a wrapper class for tracing
class ModelWrapper(torch.nn.Module):
def __init__(self, model):
super().__init__()
self.model = model
def forward(self, input_ids):
return self.model(input_ids).logits
# Wrap and trace model
wrapped_model = ModelWrapper(model)
traced_model = torch.jit.trace(wrapped_model, inputs.input_ids)
# Convert to CoreML
model_mlpackage = ct.convert(
traced_model,
inputs=[
ct.TensorType(
name="input_ids",
shape=inputs.input_ids.shape,
dtype=np.int32
)
],
source="pytorch",
minimum_deployment_target=ct.target.iOS16,
convert_to="mlprogram",
compute_precision=ct.precision.FLOAT16,
compute_units=ct.ComputeUnit.ALL,
)
model_mlpackage.save("LLaMmlein_1B.mlpackage")
|