LLaMmlein_1B_CoreML / convert_model.py
dotwee's picture
add model
4e6f878
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")