Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+tokenizer.json filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,152 @@
+---
+license: mit
+tags:
+- coreml
+- ANE
+- LLaMA
+- Qwen
+- DeepSeek
+- Apple
+- Apple Neural Engine
+- DeepHermes
+---
+# ANEMLL
+
+**ANEMLL** (pronounced like "animal") is an open-source project focused on accelerating the porting of Large Language Models (LLMs) to tensor processors, starting with the Apple Neural Engine (ANE).
+
+The goal is to provide a fully open-source pipeline from model conversion to inference for common LLM architectures running on ANE.
+
+This enables seamless integration and on-device inference for low-power applications on edge devices, ensuring maximum privacy and security.
+
+This is critical for autonomous applications, where models run directly on the device without requiring an internet connection.
+
+For more information, visit the [ANEMLL GitHub repository](https://github.com/anemll/anemll).
+
+
+---
+
+## License
+
+ANEMLL is licensed under the [MIT License](https://opensource.org/license/mit).  
+The original model may require a separate license depending on the architecture:
+- LLaMA models: Based on Meta's LLaMA and may require Meta's license
+- Qwen models: Based on Alibaba's Qwen and may require Alibaba's license
+- Other models: Check respective original model licenses
+
+This model is converted for CoreML using ANEMLL's open-source conversion pipeline. It supports multiple LLM architectures including LLaMA, Qwen, and DeepSeek variants.
+
+---
+
+## Requirements
+
+- **macOS Sequoia** with Apple Neural Engine and 8GB RAM or more
+- **CoreML Tools** and **HuggingFace Transformers** libraries 
+- **Python 3.9**
+
+`chat.py` provides a sample inference script.  
+`chat_full.py` provides a sample inference script with history and conversation management.  
+
+**Installation**
+
+1. Download the model from Hugging Face:
+```bash
+# Install required tools
+pip install huggingface_hub
+
+# Install Git LFS (Large File Support)
+# macOS with Homebrew:
+brew install git-lfs
+# Or Ubuntu/Debian:
+# sudo apt-get install git-lfs
+
+# Initialize Git LFS
+git lfs install
+
+# Clone the repository with model files
+git clone https://huggingface.co/anemll/anemll-Qwen3-4B-ctx1024_0.3.0
+```
+
+2. Extract model files:
+```bash
+# Navigate to cloned directory
+cd anemll-Qwen3-4B-ctx1024_0.3.0
+
+# Pull LFS files (model weights)
+git lfs pull
+
+# Extract CoreML model files
+find . -type f -name "*.zip" -exec unzip {} \;
+```
+
+3. Install dependencies:
+```bash
+pip install coremltools transformers
+```
+
+**Coremltools:**
+
+See coremltools installation guide at https://coremltools.readme.io/v4.0/docs/installation 
+
+**How to Run**
+
+1. Basic chat interface:
+```bash
+python chat.py --meta ./meta.yaml
+```
+
+2. Full conversation mode with history:
+```bash
+python chat_full.py --meta ./meta.yaml
+```
+
+> Note: The first time the model loads, macOS will take some time to place it on the device.
+> Subsequent loads will be instantaneous.
+> Use Ctrl-D to exit, Ctrl-C to interrupt inference.
+
+**More Info**
+Please check following links for later updates:
+
+* [GitHub](https://github.com/anemll)
+* [Hugging Face Models](https://huggingface.co/anemll)
+* [Twitter/X](https://x.com/anemll)
+* [Website](https://anemll.com)
+
+
+[email protected]
+
+# anemll-Qwen3-4B-ctx1024_0.3.0
+
+This is a CoreML model converted using ANEMLL for Apple Neural Engine inference.
+
+## Available Distributions
+
+### Standard Distribution
+- Contains zipped MLMODELC files
+- Suitable for macOS and development
+
+### iOS Distribution
+- Contains unzipped MLMODELC files
+- Ready for iOS deployment
+- Includes offline tokenizer support
+
+## Model Information
+- Context Length: 1024
+- Batch Size: 64
+- Number of Chunks: 2
+
+## Quick Start
+
+### Test in iOS/macOS App
+Try our sample Chat-Bot app on TestFlight:
+1. Install TestFlight from App Store
+2. Join beta test: [TestFlight Link](https://testflight.apple.com/join/jrQq1D1C)
+3. App includes a small demo model pre-installed
+4. You can add custom models via HuggingFace URLs
+
+> [!Note]
+> - The TestFlight app works on both iOS and macOS
+> - Demonstrates proper model integration and provides a reference implementation
+> - iOS requires unzipped MLMODELC files and config.json for offline tokenizer
+> - macOS supports both zipped and unzipped model formats
+
+```

chat.py

@@ -0,0 +1,949 @@
+# chat.py
+#!/usr/bin/env python3
+# chat.py
+# Copyright (c) 2025 Anemll
+# Licensed under the MIT License
+
+import argparse
+import os
+import re
+import glob
+from pathlib import Path
+import coremltools as ct
+from transformers import LlamaTokenizer, AutoTokenizer
+import torch
+import torch.nn.functional as F
+import numpy as np
+import queue
+import threading
+import time
+import yaml
+import sys
+
+# ANSI color codes
+LIGHT_BLUE = "\033[94m"
+DARK_BLUE = "\033[34m"
+LIGHT_GREEN = "\033[92m"
+RESET_COLOR = "\033[0m"
+
+# Add at top with other constants
+WARMUP_TOKEN_LIMIT = 10  # Maximum tokens to generate during warmup
+
+class TokenPrinter:
+    """Handles background printing of generated tokens."""
+    def __init__(self, tokenizer):
+        self.tokenizer = tokenizer
+        self.token_queue = queue.Queue()
+        self.stop_event = threading.Event()
+        self.thread = None
+        self.buffer = ""
+        self.lock = threading.Lock()
+        self.thinking = True  # Track if we're still in thinking mode
+        self.decoding_buffer = []  # Buffer for token IDs
+        # Add token counting and timing
+        self.start_time = time.time()
+        self.token_count = 0
+        self.start()
+
+    def start(self):
+        """Start the printer thread."""
+        if self.thread is None:
+            self.thread = threading.Thread(target=self._print_worker)
+            self.thread.daemon = True
+            self.thread.start()
+
+    def add_token(self, token_id):
+        """Add a token to the print queue."""
+        if not self.stop_event.is_set():
+            self.token_queue.put(token_id)
+            self.token_count += 1
+
+    def drain_buffer(self):
+        """Decode token IDs from decoding_buffer in the main thread."""
+        if not self.decoding_buffer:
+            return
+
+        # Decode all tokens at once in the main thread
+        token_str = self.tokenizer.decode(self.decoding_buffer)
+        self.decoding_buffer.clear()
+        
+        # Store the text in buffer for later saving to file
+        with self.lock:
+            self.buffer += token_str
+
+        # Color-handling logic
+        if self.thinking and "</think>" in token_str:
+            self.thinking = False
+            parts = token_str.split("</think>")
+            if len(parts) > 0:
+                print(parts[0] + "</think>", end='', flush=True)
+                if len(parts) > 1:
+                    print(LIGHT_BLUE + parts[1], end='', flush=True)
+        else:
+            if not self.thinking:
+                print(LIGHT_BLUE + token_str, end='', flush=True)
+            else:
+                print(token_str, end='', flush=True)
+
+    def _print_worker(self):
+        """Worker thread that takes token_ids from the queue."""
+        while not self.stop_event.is_set():
+            try:
+                token_id = self.token_queue.get(timeout=0.01)
+                with self.lock:
+                    self.decoding_buffer.append(token_id)
+                self.token_queue.task_done()
+            except queue.Empty:
+                continue
+            except Exception as e:
+                print(f"\nError: Token printer error: {str(e)}")
+                break
+
+    def stop(self):
+        """Stop the printer thread."""
+        if self.thread and self.thread.is_alive():
+            # Ensure any remaining tokens are processed
+            self.drain_buffer()
+            self.stop_event.set()
+            try:
+                self.thread.join(timeout=1.0)
+            except Exception:
+                pass
+            # Calculate and print tokens/s with shorter format in blue
+            elapsed = time.time() - self.start_time
+            if elapsed > 0 and self.token_count > 0:
+                tokens_per_sec = self.token_count / elapsed
+                print(f"\n{DARK_BLUE}{tokens_per_sec:.1f} t/s{RESET_COLOR}")
+            else:
+                print(RESET_COLOR)  # Reset color at the end
+        return self.buffer
+
+def parse_model_path(path):
+    """Parse model path and return full path with .mlmodelc or .mlpackage extension."""
+    path = Path(path)
+    
+    # If path exists exactly as specified, return it
+    if path.exists():
+        return str(path)
+        
+    # Try with both extensions
+    candidates = [
+        path,  # Original path
+        path.with_suffix('.mlmodelc'),  # With .mlmodelc
+        path.with_suffix('.mlpackage'),  # With .mlpackage
+        Path(str(path) + '.mlmodelc'),  # Handle case where extension is included
+        Path(str(path) + '.mlpackage')
+    ]
+    
+    # Try all possible paths
+    for candidate in candidates:
+        if candidate.exists():
+            print(f"Found model at: {candidate}")
+            return str(candidate)
+            
+    # If we get here, no valid path was found
+    print("\nError: Model not found. Tried following paths:")
+    for candidate in candidates:
+        print(f"  {candidate}")
+    raise FileNotFoundError(f"Model not found: {path}")
+
+def parse_ffn_filename(path):
+    """Parse FFN model filename to extract chunk information."""
+    path = Path(path)
+    pattern = r'FFN_PF.*_chunk_(\d+)of(\d+)'
+    match = re.search(pattern, path.name)
+    
+    if match:
+        current_chunk = int(match.group(1))
+        total_chunks = int(match.group(2))
+        return current_chunk, total_chunks
+    return None, None
+
+def find_all_chunks(base_path):
+    """Find all chunk files matching the base FFN path pattern."""
+    path = Path(base_path)
+    pattern = re.sub(r'_chunk_\d+of\d+', '_chunk_*', str(path))
+    return sorted(glob.glob(pattern))
+
+def load_model(path, function_name=None):
+    """Load a CoreML model, handling both .mlmodelc and .mlpackage formats."""
+    path = Path(path)
+    compute_unit = ct.ComputeUnit.CPU_AND_NE
+    
+    try:
+        if path.suffix == '.mlmodelc':
+            # For compiled models (.mlmodelc), use CompiledMLModel
+            if function_name:
+                return ct.models.CompiledMLModel(str(path), compute_unit, function_name=function_name)
+            else:
+                return ct.models.CompiledMLModel(str(path), compute_unit)
+        else:
+            # For packages (.mlpackage)
+            if function_name:
+                return ct.models.MLModel(str(path), function_name=function_name)
+            else:
+                return ct.models.MLModel(str(path))
+                
+    except RuntimeError as e:
+        if "valid manifest does not exist" in str(e):
+            print(f"\nError: Could not load compiled model at {path}")
+            print("This might be because:")
+            print("1. The model is not properly compiled")
+            print("2. The model was compiled for a different OS version")
+            print("3. The model needs to be recompiled")
+            print("\nTry using the .mlpackage version instead, or recompile the model.")
+        raise
+
+def load_metadata(model,args):
+    # Extract metadata and config parameters
+    metadata = {}
+    if hasattr(model, 'user_defined_metadata'):
+        meta = model.user_defined_metadata
+        
+        # Extract key parameters with defaults
+        metadata['context_length'] = int(meta.get('com.anemll.context_length', 512))
+        metadata['state_length'] = int(meta.get('com.anemll.state_length', metadata['context_length']))  # Added state_length
+        metadata['batch_size'] = int(meta.get('com.anemll.batch_size', 64))
+        metadata['lut_bits'] = int(meta.get('com.anemll.lut_bits', 0))
+        metadata['num_chunks'] = int(meta.get('com.anemll.num_chunks', 1))
+        
+        print("\nExtracted Parameters:")
+        print(f"  Context Length: {metadata['context_length']}")
+        print(f"  State Length: {metadata['state_length']}")
+        print(f"  Prefill Batch Size: {metadata['batch_size']}")
+        print(f"  LUT Bits: {metadata['lut_bits']}")
+        print(f"  Number of Chunks: {metadata['num_chunks']}")
+        
+        # Print model info
+        print("\nModel Info:")
+        if 'com.anemll.info' in meta:
+            print(f"  {meta['com.anemll.info']}")
+        if 'com.github.apple.coremltools.version' in meta:
+            print(f"  CoreML Tools: {meta['com.github.apple.coremltools.version']}")
+        
+        # Print model input/output shapes
+        print("\nModel Shapes:")
+        if hasattr(model, 'input_description'):
+            print("  Inputs:")
+            for name, desc in model.input_description.items():
+                print(f"    {name}: {desc}")
+        if hasattr(model, 'output_description'):
+            print("  Outputs:")
+            for name, desc in model.output_description.items():
+                print(f"    {name}: {desc}")
+    else:
+        print("\nWarning: No metadata found in model")
+
+        # Check if model directory name contains context length pattern (ctxXXX)
+        ctx_len = 512
+        if args.context_length is  None:
+            import re
+            ctx_match = re.search(r'ctx(\d+)', str(args.d))
+            if ctx_match:
+                ctx_len0 = int(ctx_match.group(1))
+                if 512 <= ctx_len0 <= 8096:
+                    ctx_len = ctx_len0
+                    print(f"\nDetected context length {ctx_len} from directory name")
+            else:
+                print(f"\nWarning: No context length found in directory  {ctx_len} from directory name {args.d}")
+        else:
+            ctx_len = args.context_length
+
+        # Use defaults or values from args
+        metadata['context_length'] = ctx_len
+        metadata['state_length'] = ctx_len
+        # Get batch size from args or use default
+        metadata['batch_size'] = getattr(args, 'batch_size', 64)
+        metadata['lut_bits'] = 4
+        metadata['num_chunks'] = getattr(args, 'num_chunks', 4)
+        print("\nUsing parameters:")
+        print(f"  Context Length: {metadata['context_length']}")
+        print(f"  State Length: {metadata['state_length']}")
+        print(f"  Prefill Batch Size: {metadata['batch_size']}")
+        print(f"  LUT Bits: {metadata['lut_bits']}")
+        print(f"  Number of Chunks: {metadata['num_chunks']}")
+
+    # Override with values from args if they exist
+    if hasattr(args, 'batch_size') and args.batch_size is not None:
+        metadata['batch_size'] = args.batch_size
+        print(f"\nOverriding batch size from args: {args.batch_size}")
+    if hasattr(args, 'num_chunks') and args.num_chunks is not None:
+        metadata['num_chunks'] = args.num_chunks
+        print(f"\nOverriding num chunks from args: {args.num_chunks}")
+    
+    return metadata
+    
+def load_models(args,metadata):
+    """Load all required models and extract metadata."""
+    print("\nLoading models...")
+    
+    try:
+        # Load embeddings model
+        print("\nLoading embeddings model...")
+        embed_path = parse_model_path(args.embed)
+        print(f"Loading from: {embed_path}")
+        embed_model = load_model(embed_path)
+        print("Embeddings model loaded successfully")
+        metadata = load_metadata(embed_model,args)
+        
+
+        
+        # Load LM head model
+        print("\nLoading LM head model...")
+        lmhead_path = parse_model_path(args.lmhead)
+        print(f"Loading from: {lmhead_path}")
+        lmhead_model = load_model(lmhead_path)
+        print("LM head model loaded successfully")
+        
+        # Parse FFN path and find chunks if needed
+        print("\nLoading FFN+PREFILL model(s)...")
+        ffn_path = parse_model_path(args.ffn)
+        chunk_no, total_chunks = parse_ffn_filename(ffn_path)
+        
+        ffn_models = []
+        if chunk_no and total_chunks:
+            print(f"\nDetected chunked FFN+PREFILL model ({total_chunks} chunks)")
+            # Find and load all chunks
+            chunk_paths = find_all_chunks(ffn_path)
+            if len(chunk_paths) != total_chunks:
+                raise ValueError(f"Found {len(chunk_paths)} chunks but filename indicates {total_chunks} chunks")
+                
+            for chunk_path in chunk_paths:
+                print(f"\nLoading FFN+PREFILL chunk: {Path(chunk_path).name}")
+                try:
+                    # For chunked models, we need both infer and prefill functions
+                    ffn_models.append({
+                        'infer': load_model(chunk_path, function_name='infer'),
+                        'prefill': load_model(chunk_path, function_name='prefill')
+                    })
+                    print("Chunk loaded successfully")
+                except Exception as e:
+                    print(f"Error loading chunk {chunk_path}: {str(e)}")
+                    raise
+            metadata = load_metadata(ffn_models[0],args)
+
+        else:
+            print("\nLoading single FFN model...")
+            ffn_models.append(load_model(ffn_path))
+            print("FFN model loaded successfully")
+        
+        return embed_model, ffn_models, lmhead_model, metadata
+        
+    except Exception as e:
+        print(f"\nError loading models: {str(e)}")
+        print("\nPlease ensure all model files exist and are accessible.")
+        print("Expected files:")
+        print(f"  Embeddings: {args.embed}")
+        print(f"  LM Head: {args.lmhead}")
+        print(f"  FFN: {args.ffn}")
+        raise
+
+# At the top of the file, make this a default path
+
+def initialize_tokenizer(model_path=None):
+    """Initialize and configure the tokenizer."""
+    try:
+
+        
+        tokenizer = AutoTokenizer.from_pretrained(
+            str(model_path), 
+            use_fast=False,
+            trust_remote_code=True
+        )
+        
+        print("\nTokenizer Configuration:")
+        print(f"Tokenizer type: {type(tokenizer)}")
+        print(f"Tokenizer name: {tokenizer.__class__.__name__}")
+        print(f"Vocabulary size: {len(tokenizer)}")
+        print(f"Model max length: {tokenizer.model_max_length}")
+
+        if tokenizer.pad_token is None:
+            tokenizer.pad_token = tokenizer.eos_token
+            tokenizer.pad_token_id = tokenizer.eos_token_id
+            print("Set PAD token to EOS token")
+        
+        tokenizer.padding_side = "left"
+        
+        print(f"\nSpecial Tokens:")
+        print(f"PAD token: '{tokenizer.pad_token}' (ID: {tokenizer.pad_token_id})")
+        print(f"EOS token: '{tokenizer.eos_token}' (ID: {tokenizer.eos_token_id})")
+        print(f"BOS token: '{tokenizer.bos_token}' (ID: {tokenizer.bos_token_id})")
+        print(f"UNK token: '{tokenizer.unk_token}' (ID: {tokenizer.unk_token_id})")
+
+        return tokenizer
+        
+    except Exception as e:
+        print(f"\nError: Failed to load tokenizer from {model_path}")
+        print(f"Error details: {str(e)}")
+        print(f"Error type: {type(e)}")
+        print("\nThis appears to be a tokenizer loading issue.")
+        
+        # Check if it's the specific Qwen tokenizer file issue
+        if "expected str, bytes or os.PathLike object, not NoneType" in str(e):
+            print("\nThis error suggests the tokenizer files are missing or incomplete.")
+            print("For Qwen models, you need the original model directory with tokenizer files.")
+            print("Try using: --tokenizer ~/.cache/huggingface/hub/models--Qwen--Qwen3-0.6B/snapshots/YOUR_SNAPSHOT_ID")
+        else:
+            print("Please provide the path to a compatible model directory with tokenizer files.")
+        import traceback
+        traceback.print_exc()
+        raise
+
+
+
+def make_causal_mask(length, start):
+    """Create causal attention mask."""
+    mask = np.full((1, 1, length, length), -np.inf, dtype=np.float16)
+    row_indices = np.arange(length).reshape(length, 1)
+    col_indices = np.arange(length).reshape(1, length)
+    mask[:, :, col_indices <= (row_indices + start)] = 0
+    return mask
+
+def initialize_causal_mask(context_length):
+    """Initialize causal mask for transformer attention."""
+    causal_mask = make_causal_mask(context_length, 0)
+    causal_mask = torch.tensor(causal_mask, dtype=torch.float16)
+    print(f"\nInitialized causal mask for context length {context_length}")
+    return causal_mask
+
+def run_prefill(embed_model, ffn_models, input_ids, context_pos, context_length, batch_size=64, state=None, causal_mask=None):
+    """Run prefill on the input sequence."""
+    # Use provided causal mask or create one if not provided
+    if causal_mask is None:
+        causal_mask = make_causal_mask(context_length, 0)
+        causal_mask = torch.tensor(causal_mask, dtype=torch.float16)
+    
+    # Process in batches
+    batch_pos = 0
+    while batch_pos < context_pos:
+        batch_end = min(batch_pos + batch_size, context_pos)
+        current_batch_size = batch_end - batch_pos
+        
+        # Get current batch
+        batch_input = input_ids[:, batch_pos:batch_end]
+        
+        # Always pad to full batch size for prefill
+        batch_input = F.pad(
+            batch_input,
+            (0, batch_size - current_batch_size),
+            value=0
+        )
+        
+        # Generate position IDs for full batch size
+        position_ids = torch.arange(batch_pos, batch_pos+batch_size, dtype=torch.int32)  # Changed: Always use full batch size
+        batch_causal_mask = causal_mask[:, :, batch_pos:batch_pos+batch_size, :]  # Changed: Use full batch size
+        
+        # Run embeddings with proper batch size
+        hidden_states = torch.from_numpy(
+            embed_model.predict({
+                'input_ids': batch_input.numpy(),
+                'batch_size': np.array([batch_size], dtype=np.int32)  # Add batch_size parameter
+            })['hidden_states']
+        )
+        
+        # Run through FFN chunks with state
+        for ffn_model in ffn_models:
+            if isinstance(ffn_model, dict):
+                inputs = {
+                    'hidden_states': hidden_states.numpy(),  # [1, 64, hidden_size]
+                    'position_ids': position_ids.numpy(),    # [64]
+                    'causal_mask': batch_causal_mask.numpy(), # [1, 1, 64, context_length]
+                    'current_pos': np.array([batch_pos], dtype=np.int32)  # [1]
+                }
+                output = ffn_model['prefill'].predict(inputs, state)
+                hidden_states = torch.from_numpy(output['output_hidden_states'])
+        
+        batch_pos = batch_end
+    
+    return torch.tensor([context_pos], dtype=torch.int32)
+
+def generate_next_token(embed_model, ffn_models, lmhead_model, input_ids, pos, context_length, metadata, state=None, causal_mask=None, temperature=0.0):
+    """Generate the next token."""
+    # Get current token
+    current_token = input_ids[:, pos-1:pos]  # [1, 1]
+    
+    # Run embeddings
+    hidden_states = torch.from_numpy(
+        embed_model.predict({'input_ids': current_token.numpy()})['hidden_states']
+    )  # [1, 1, hidden_size]
+    
+    # Create masks
+    update_mask = torch.zeros((1, 1, context_length, 1), dtype=torch.float16)
+    update_mask[0, 0, pos-1, 0] = 1.0
+    position_ids = torch.tensor([pos-1], dtype=torch.int32)  # [1]
+    
+    # Use provided causal mask or create one if not provided
+    if causal_mask is None:
+        causal_mask_data = make_causal_mask(context_length, 0)
+        single_causal_mask = torch.tensor(causal_mask_data[:, :, pos-1:pos, :], dtype=torch.float16)  # [1, 1, 1, context_length]
+    else:
+        single_causal_mask = causal_mask[:, :, pos-1:pos, :]
+    
+    # Run through FFN chunks with state
+    for ffn_model in ffn_models:
+        if isinstance(ffn_model, dict):
+            inputs = {
+                'hidden_states': hidden_states.numpy(),
+                'update_mask': update_mask.numpy(),
+                'position_ids': position_ids.numpy(),
+                'causal_mask': single_causal_mask.numpy(),
+                'current_pos': position_ids.numpy()
+            }
+            output = ffn_model['infer'].predict(inputs, state)
+            hidden_states = torch.from_numpy(output['output_hidden_states'])
+    
+    # Run LM head
+    lm_output = lmhead_model.predict({'hidden_states': hidden_states.numpy()})
+    # Debug print
+    #print("\nLM Head output keys:", list(lm_output.keys()))
+    
+    # Get number of logits from metadata, using split_lm_head if available
+    # First check for split_lm_head (new), then num_logits (legacy), default to 8
+    num_logits = metadata.get('split_lm_head', metadata.get('num_logits', 8))
+    
+    # Combine logits1-N if they exist
+    if 'logits1' in lm_output:
+        # Concatenate all logits parts
+        logits_parts = []
+        for i in range(1, num_logits + 1):
+            key = f'logits{i}'
+            if key in lm_output:
+                logits_parts.append(torch.from_numpy(lm_output[key]))
+        logits = torch.cat(logits_parts, dim=-1)  # Concatenate along vocab dimension
+    else:
+        # Try output_logits as fallback
+        logits = torch.from_numpy(lm_output['output_logits'])
+    
+    # Apply temperature and sample
+    if temperature > 0:
+        logits = logits / temperature
+        probs = F.softmax(logits[0, -1, :], dim=-1)
+        next_token = torch.multinomial(probs, num_samples=1).item()
+    else:
+        next_token = torch.argmax(logits[0, -1, :]).item()
+    
+    return next_token
+
+def create_unified_state(ffn_models, context_length):
+    """Create unified KV cache state for transformer."""
+    if isinstance(ffn_models[0], dict):
+        # Use first FFN model's prefill function to create state
+        state = ffn_models[0]['prefill'].make_state()
+        print(f"\nCreated unified transformer state for {len(ffn_models)} chunks")
+        return state
+    else:
+        state = ffn_models[0].make_state()
+        print("\nCreated unified transformer state")
+        return state
+
+def chat_loop(embed_model, ffn_models, lmhead_model, tokenizer, metadata, state, causal_mask=None, auto_prompt=None, warmup=False, save_file=None):
+    """Interactive chat loop."""
+    context_length = metadata.get('context_length')
+    batch_size = metadata.get('batch_size', 64)
+    
+    if not warmup:
+        print(f"\nUsing context length: {context_length}")
+        print("\nStarting chat session. Press Ctrl+D to exit.")
+        print("Type your message and press Enter to chat.")
+    
+    # Check if tokenizer has chat template and if it works
+    has_chat_template = False
+    try:
+        # Test if chat template works
+        test_messages = [{"role": "user", "content": "test"}]
+        tokenizer.apply_chat_template(test_messages, return_tensors="pt")
+        has_chat_template = True
+        if not warmup:
+            print("\nUsing chat template for prompts")
+    except:
+        if not warmup:
+            print("\nUsing manual formatting for prompts")
+    
+    conversation = []
+    
+    try:
+        while True:
+            try:
+                if not warmup:
+                    print(f"\n{LIGHT_GREEN}You:{RESET_COLOR}", end=' ', flush=True)
+                if auto_prompt is not None:
+                    user_input = auto_prompt
+                    if not warmup:
+                        print(user_input)
+                else:
+                    user_input = input().strip()
+            except EOFError:
+                if not warmup:
+                    print("\nExiting chat...")
+                break
+                
+            if not user_input:
+                continue
+            
+            # Format prompt based on tokenizer capabilities
+            if has_chat_template:
+                messages = [{"role": "user", "content": user_input}]
+                input_ids = tokenizer.apply_chat_template(
+                    messages,
+                    return_tensors="pt",
+                    add_generation_prompt=True
+                ).to(torch.int32)
+            else:
+                # Manual formatting for Llama models without chat template
+                formatted_prompt = f"[INST] {user_input} [/INST]"
+                input_ids = tokenizer(
+                    formatted_prompt,
+                    return_tensors="pt",
+                    add_special_tokens=True
+                ).input_ids.to(torch.int32)
+            
+            context_pos = input_ids.size(1)
+            
+            if not warmup:
+                print(f"\n{LIGHT_BLUE}Assistant:{RESET_COLOR}", end=' ', flush=True)
+            
+            # Initialize token printer
+            token_printer = TokenPrinter(tokenizer)
+            tokens_generated = 0  # Track number of tokens
+            
+            try:
+                # Start prefill timing
+                prefill_start = time.time()
+                
+                # Run prefill with state and causal mask
+                current_pos = run_prefill(
+                    embed_model,
+                    ffn_models,
+                    input_ids,
+                    context_pos,
+                    context_length,
+                    batch_size,
+                    state,
+                    causal_mask
+                )
+                
+                # Calculate prefill timing
+                prefill_time = time.time() - prefill_start
+                prefill_tokens = context_pos  # Number of tokens in input
+                prefill_tokens_per_sec = prefill_tokens / prefill_time if prefill_time > 0 else 0
+                
+                # Generation loop with state
+                input_ids = input_ids
+                pos = context_pos
+                inference_start = time.time()
+                inference_tokens = 0
+                
+                while pos < context_length - 1:
+                    # Generate next token with causal mask
+                    next_token = generate_next_token(
+                        embed_model,
+                        ffn_models,
+                        lmhead_model,
+                        input_ids,
+                        pos,
+                        context_length,
+                        metadata,
+                        state,
+                        causal_mask
+                    )
+                    
+                    # Add token to sequence
+                    if pos < input_ids.size(1):
+                        input_ids[0, pos] = next_token
+                    else:
+                        input_ids = torch.cat([
+                            input_ids,
+                            torch.tensor([[next_token]], dtype=torch.int32)
+                        ], dim=1)
+                    
+                    # Add to printer only if not in warmup
+                    if not warmup:
+                        token_printer.add_token(next_token)
+                        token_printer.drain_buffer()
+                    
+                    pos += 1
+                    tokens_generated += 1
+                    inference_tokens += 1
+                    
+                    # Check limits
+                    if warmup and tokens_generated >= WARMUP_TOKEN_LIMIT:
+                        break
+                        
+                    if next_token == tokenizer.eos_token_id:
+                        break
+                
+                # Calculate inference timing
+                inference_time = time.time() - inference_start
+                inference_tokens_per_sec = inference_tokens / inference_time if inference_time > 0 else 0
+                
+                # Get final response and add to conversation
+                if not warmup:
+                    response = token_printer.stop()
+                    # Print timing stats
+                    prefill_ms = prefill_time * 1000  # Convert to milliseconds
+                    print(f"\nPrefill: {prefill_ms:.1f}ms ({prefill_tokens_per_sec:.1f} t/s)")
+                    print(f"Inference: {inference_tokens_per_sec:.1f} t/s")
+                    print(f"Total: Generated {tokens_generated} tokens in {prefill_time + inference_time:.2f}s")
+                    conversation.append({"role": "assistant", "content": response})
+                    
+                    # Save response to file if requested
+                    if save_file:
+                        try:
+                            # Add small delay to ensure all tokens are processed
+                            time.sleep(0.5)
+                            
+                            # Make sure response ends with EOS token if it's supposed to
+                            if response and not response.endswith("<|eot_id|>") and not response.endswith("</s>"):
+                                if tokenizer.eos_token:
+                                    eos_text = tokenizer.decode([tokenizer.eos_token_id])
+                                    if not response.endswith(eos_text):
+                                        print(f"\n{DARK_BLUE}Adding missing EOS token for consistency{RESET_COLOR}")
+                                        response += eos_text
+                            
+                            with open(save_file, 'w') as f:
+                                f.write(response)
+                            print(f"\n{DARK_BLUE}Response saved to file: {save_file}{RESET_COLOR}")
+                        except Exception as e:
+                            print(f"\n{DARK_BLUE}Error saving to file: {str(e)}{RESET_COLOR}")
+                else:
+                    token_printer.stop()  # Clean up without printing stats
+                
+                # Exit after one response in auto_prompt mode
+                if auto_prompt is not None:
+                    break
+                
+            except KeyboardInterrupt:
+                print("\nGeneration interrupted")
+                token_printer.stop()
+                continue
+                
+    except Exception as e:
+        print(f"\nError in chat loop: {str(e)}")
+        import traceback
+        traceback.print_exc()
+
+def parse_args():
+    parser = argparse.ArgumentParser(description='Chat with CoreML LLaMA, gil resolved  (c) 2025 Anemll')
+    
+    # Add meta.yaml option
+    parser.add_argument('--meta', type=str, help='Path to meta.yaml to load all parameters')
+    
+    # Model paths
+    parser.add_argument('--d', '--dir', type=str, default='.',
+                       help='Directory containing model files (default: current directory)')
+    parser.add_argument('--embed', type=str, required=False,
+                       help='Path to embeddings model (relative to --dir)')
+    parser.add_argument('--ffn', type=str, required=False,
+                       help='Path to FFN model (can be chunked, relative to --dir)')
+    parser.add_argument('--lmhead', type=str, required=False,
+                       help='Path to LM head model (relative to --dir)')
+    parser.add_argument('--tokenizer', type=str, required=False,
+                       help='Path to tokenizer')
+    
+    # Add new argument for auto-generation
+    parser.add_argument('--prompt', type=str,
+                       help='If specified, run once with this prompt and exit')
+    
+    # Add save option
+    parser.add_argument('--save', type=str,
+                       help='Save assistant\'s response to specified file')
+    
+    # Add no-warmup flag
+    parser.add_argument('--nw', action='store_true',
+                       help='Skip warmup phase')
+    
+    # Model configuration
+    parser.add_argument('--context-length', type=int,
+                       help='Context length for the model (default: 512), if not provided, it will be detected from the model directory name ctxNUMBER')
+    parser.add_argument('--batch-size', type=int,
+                       help='Batch size for prefill (default: 64)')
+    parser.add_argument('--num-logits', type=int, default=8,
+                       help='Number of logits outputs from LM head (default: 8, legacy)')
+    parser.add_argument('--split-lm-head', type=int, 
+                       help='Number of logits splits from LM head (default: 8 for llama, 16 for qwen)')
+    
+    args = parser.parse_args()
+    
+    # If meta.yaml is provided, load parameters from it
+    if args.meta:
+        try:
+            with open(args.meta, 'r') as f:
+                meta = yaml.safe_load(f)
+            params = meta['model_info']['parameters']
+            
+            # Set model directory to meta.yaml directory if not specified
+            if not args.d or args.d == '.':
+                args.d = str(Path(args.meta).parent)
+            
+            # Build model paths based on parameters
+            prefix = params.get('model_prefix', 'llama')  # Default to 'llama' if not specified
+            lut_ffn = f"_lut{params['lut_ffn']}" if params['lut_ffn'] != 'none' else ''
+            lut_lmhead = f"_lut{params['lut_lmhead']}" if params['lut_lmhead'] != 'none' else ''
+            lut_embeddings = f"_lut{params['lut_embeddings']}" if params['lut_embeddings'] != 'none' else ''
+            num_chunks = int(params['num_chunks'])
+            
+            # Set model paths if not specified
+            if not args.lmhead:
+                args.lmhead = f'{prefix}_lm_head{lut_lmhead}'
+            if not args.embed:
+                args.embed = f'{prefix}_embeddings{lut_embeddings}'  # Changed from lm_head to embeddings
+            if not args.ffn:
+                args.ffn = f'{prefix}_FFN_PF{lut_ffn}_chunk_01of{num_chunks:02d}'
+            if not args.tokenizer:
+                # Check if there's a tokenizer_path parameter in meta.yaml
+                if 'tokenizer_path' in params:
+                    args.tokenizer = params['tokenizer_path']
+                else:
+                    # Default to the model directory, but this might need manual override
+                    args.tokenizer = args.d
+            
+            # Set other parameters if not overridden by command line
+            if args.context_length is None:
+                args.context_length = int(params['context_length'])
+            if args.batch_size is None:
+                args.batch_size = int(params['batch_size'])
+            args.num_chunks = num_chunks
+            # Add num_logits parameter with default of 8, override command line if present in meta
+            if 'num_logits' in params:
+                args.num_logits = int(params['num_logits'])
+            
+            # Add split_lm_head parameter with default of 8
+            if 'split_lm_head' in params:
+                args.split_lm_head = int(params['split_lm_head'])
+            else:
+                args.split_lm_head = 8  # Default value for backward compatibility
+            
+            print(f"\nLoaded parameters from {args.meta}:")
+            print(f"  Context Length: {args.context_length}")
+            print(f"  Batch Size: {args.batch_size}")
+            print(f"  Num Chunks: {args.num_chunks}")
+            print(f"  Num Logits: {args.num_logits}")
+            print(f"  Split LM Head: {args.split_lm_head}")
+            print(f"  Models Directory: {args.d}")
+            print(f"  Embeddings: {args.embed}")
+            print(f"  LM Head: {args.lmhead}")
+            print(f"  FFN: {args.ffn}")
+            
+        except Exception as e:
+            print(f"\nError loading meta.yaml: {str(e)}")
+            sys.exit(1)
+    else:
+        # If no meta.yaml, set default split_lm_head if not provided
+        if not hasattr(args, 'split_lm_head') or args.split_lm_head is None:
+            args.split_lm_head = args.num_logits  # Use num_logits as fallback
+    
+    return args
+
+def main():
+    args = parse_args()
+    
+    # Convert directory to absolute path
+    model_dir = Path(args.d).resolve()
+    if not model_dir.exists():
+        print(f"\nError: Model directory not found: {model_dir}")
+        return 1
+        
+    print(f"\nUsing model directory: {model_dir}")
+    print(f"Context length: {args.context_length}")
+    
+    try:
+        # Update paths to be relative to model directory
+        args.embed = str(model_dir / args.embed)
+        args.ffn = str(model_dir / args.ffn)
+        args.lmhead = str(model_dir / args.lmhead)
+        
+        # Handle tokenizer path separately since it's not relative to model_dir
+        if args.tokenizer is None:
+            args.tokenizer = str(model_dir)
+        
+        # Check if tokenizer directory exists and has required files
+        tokenizer_path = Path(args.tokenizer)
+        if not tokenizer_path.exists():
+            print(f"\nError: Tokenizer directory not found: {args.tokenizer}")
+            return 1
+        
+        # Check if tokenizer has the required files
+        required_files = ['tokenizer.json', 'tokenizer_config.json']
+        missing_files = [f for f in required_files if not (tokenizer_path / f).exists()]
+        
+        if missing_files:
+            print(f"\nWarning: Tokenizer directory missing required files: {missing_files}")
+            print(f"Current tokenizer path: {args.tokenizer}")
+            print("\nFor Qwen models, you may need to specify the original model directory:")
+            print("  python chat.py --meta /tmp/qwen/meta.yaml --tokenizer ~/.cache/huggingface/hub/models--Qwen--Qwen3-0.6B/snapshots/YOUR_SNAPSHOT_ID")
+            print("\nOr add 'tokenizer_path' to your meta.yaml file.")
+    
+        args.tokenizer = str(Path(args.tokenizer).resolve())  # Convert to absolute path
+        print(f"Using tokenizer path: {args.tokenizer}")
+        
+        metadata = {}
+        # Load models and extract metadata
+        embed_model, ffn_models, lmhead_model, metadata = load_models(args,metadata)
+        
+        print(f"\nMetadata befor args.context_length: {metadata}")
+
+        # Override context length from command line if provided
+        if args.context_length is not None:
+            metadata['context_length'] = args.context_length
+            metadata['state_length'] = args.context_length  # Also update state_length
+            print(f"\nOverriding context length from command line: {args.context_length}")
+        
+        # Add num_logits to metadata (legacy support)
+        metadata['num_logits'] = getattr(args, 'num_logits', 8)
+        
+        # Add split_lm_head to metadata (preferred)
+        metadata['split_lm_head'] = getattr(args, 'split_lm_head', getattr(args, 'num_logits', 8))
+        
+        print(f"\nMetadata after load_models: {metadata}")
+        print(f"Using split_lm_head value: {metadata.get('split_lm_head', 8)}")
+        
+        # Load tokenizer with resolved path
+        tokenizer = initialize_tokenizer(args.tokenizer)
+        if tokenizer is None:
+            raise RuntimeError("Failed to initialize tokenizer")
+        
+        # Create unified state once
+        state = create_unified_state(ffn_models, metadata['context_length'])
+        
+        # Initialize causal mask once
+        causal_mask = initialize_causal_mask(metadata['context_length'])
+        
+        # Warmup runs to prevent Python GIL issues with CoreML !
+        if not args.nw:
+            for i in range(2):
+                chat_loop(
+                    embed_model=embed_model,
+                    ffn_models=ffn_models,
+                    lmhead_model=lmhead_model,
+                    tokenizer=tokenizer,
+                    metadata=metadata,
+                    state=state,
+                    causal_mask=causal_mask,  # Pass the causal mask
+                    warmup=True,
+                    auto_prompt="who are you?"
+                )
+        
+        # Main run
+        chat_loop(
+            embed_model=embed_model,
+            ffn_models=ffn_models,
+            lmhead_model=lmhead_model,
+            tokenizer=tokenizer,
+            metadata=metadata,
+            state=state,
+            causal_mask=causal_mask,  # Pass the causal mask
+            warmup=False,
+            auto_prompt=args.prompt,
+            save_file=args.save
+        )
+        
+    except Exception as e:
+        print(f"\nError: {str(e)}")
+        import traceback
+        traceback.print_exc()
+        return 1
+    
+    return 0
+
+if __name__ == "__main__":
+    exit(main()) 

chat_full.py

@@ -0,0 +1,989 @@
+# chat.py
+#!/usr/bin/env python3
+# chat.py
+# Copyright (c) 2025 Anemll
+# Licensed under the MIT License
+
+import argparse
+import os
+import re
+import glob
+from pathlib import Path
+import coremltools as ct
+from transformers import LlamaTokenizer, AutoTokenizer
+import torch
+import torch.nn.functional as F
+import numpy as np
+import queue
+import threading
+import time
+import yaml
+import sys
+
+# ANSI color codes
+LIGHT_BLUE = "\033[94m"
+DARK_BLUE = "\033[34m"
+LIGHT_GREEN = "\033[92m"
+RESET_COLOR = "\033[0m"
+
+# Add at the top with other constants
+WARMUP_TOKEN_LIMIT = 10  # Maximum tokens to generate during warmup
+THINKING_MODE = False
+THINKING_PROMPT = """You are a deep thinking AI, you may use extremely long chains of thought to deeply consider the problem and deliberate with yourself via systematic reasoning processes to help come to a correct solution prior to answering. You should enclose your thoughts and internal monologue inside <think> </think> tags, and then provide your solution or response to the problem."""
+DEBUG_LEVEL = 0  # Default debug level
+
+class TokenPrinter:
+    """Handles background printing of generated tokens."""
+    def __init__(self, tokenizer):
+        self.tokenizer = tokenizer
+        self.token_queue = queue.Queue()
+        self.stop_event = threading.Event()
+        self.thread = None
+        self.buffer = ""
+        self.lock = threading.Lock()
+        self.thinking = True  # Track if we're still in thinking mode
+        self.decoding_buffer = []  # Buffer for token IDs
+        # Timing and stats tracking
+        self.start_time = time.time()
+        self.token_count = 0
+        self.prefill_time = 0
+        self.inference_time = 0
+        self.context_pos = 0
+        self.start()
+
+    def start(self):
+        """Start the printer thread."""
+        if self.thread is None:
+            self.thread = threading.Thread(target=self._print_worker)
+            self.thread.daemon = True
+            self.thread.start()
+
+    def add_token(self, token_id):
+        """Add a token to the print queue."""
+        if not self.stop_event.is_set():
+            self.token_queue.put(token_id)
+            self.token_count += 1
+
+    def drain_buffer(self):
+        """Decode token IDs from decoding_buffer in the main thread."""
+        if not self.decoding_buffer:
+            return
+
+        # Decode all tokens at once in the main thread
+        token_str = self.tokenizer.decode(self.decoding_buffer)
+        self.decoding_buffer.clear()
+
+        # Color-handling logic
+        if self.thinking and "</think>" in token_str:
+            self.thinking = False
+            parts = token_str.split("</think>")
+            if len(parts) > 0:
+                print(parts[0] + "</think>", end='', flush=True)
+                if len(parts) > 1:
+                    print(LIGHT_BLUE + parts[1], end='', flush=True)
+        else:
+            if not self.thinking:
+                print(LIGHT_BLUE + token_str, end='', flush=True)
+            else:
+                print(token_str, end='', flush=True)
+
+    def _print_worker(self):
+        """Worker thread that takes token_ids from the queue."""
+        while not self.stop_event.is_set():
+            try:
+                token_id = self.token_queue.get(timeout=0.01)
+                with self.lock:
+                    self.decoding_buffer.append(token_id)
+                self.token_queue.task_done()
+            except queue.Empty:
+                continue
+            except Exception as e:
+                print(f"\nError: Token printer error: {str(e)}")
+                break
+
+    def stop(self):
+        """Stop the printer thread."""
+        if self.thread and self.thread.is_alive():
+            self.stop_event.set()
+            try:
+                self.thread.join(timeout=1.0)
+            except Exception:
+                pass
+            print(RESET_COLOR)  # Reset color at the end
+        return self.buffer
+
+    def set_timing(self, prefill_time, inference_time, context_pos):
+        """Set timing information."""
+        self.prefill_time = prefill_time
+        self.inference_time = inference_time
+        self.context_pos = context_pos
+
+def parse_model_path(path):
+    """Parse model path and return full path with .mlmodelc or .mlpackage extension."""
+    path = Path(path)
+    
+    # If path exists exactly as specified, return it
+    if path.exists():
+        return str(path)
+        
+    # Try with both extensions
+    candidates = [
+        path,  # Original path
+        path.with_suffix('.mlmodelc'),  # With .mlmodelc
+        path.with_suffix('.mlpackage'),  # With .mlpackage
+        Path(str(path) + '.mlmodelc'),  # Handle case where extension is included
+        Path(str(path) + '.mlpackage')
+    ]
+    
+    # Try all possible paths
+    for candidate in candidates:
+        if candidate.exists():
+            print(f"Found model at: {candidate}")
+            return str(candidate)
+            
+    # If we get here, no valid path was found
+    print("\nError: Model not found. Tried following paths:")
+    for candidate in candidates:
+        print(f"  {candidate}")
+    raise FileNotFoundError(f"Model not found: {path}")
+
+def parse_ffn_filename(path):
+    """Parse FFN model filename to extract chunk information."""
+    path = Path(path)
+    pattern = r'FFN_PF.*_chunk_(\d+)of(\d+)'
+    match = re.search(pattern, path.name)
+    
+    if match:
+        current_chunk = int(match.group(1))
+        total_chunks = int(match.group(2))
+        return current_chunk, total_chunks
+    return None, None
+
+def find_all_chunks(base_path):
+    """Find all chunk files matching the base FFN path pattern."""
+    path = Path(base_path)
+    pattern = re.sub(r'_chunk_\d+of\d+', '_chunk_*', str(path))
+    return sorted(glob.glob(pattern))
+
+def load_model(path, function_name=None):
+    """Load a CoreML model, handling both .mlmodelc and .mlpackage formats."""
+    path = Path(path)
+    compute_unit = ct.ComputeUnit.CPU_AND_NE
+    
+    try:
+        if path.suffix == '.mlmodelc':
+            # For compiled models (.mlmodelc), use CompiledMLModel
+            if function_name:
+                return ct.models.CompiledMLModel(str(path), compute_unit, function_name=function_name)
+            else:
+                return ct.models.CompiledMLModel(str(path), compute_unit)
+        else:
+            # For packages (.mlpackage)
+            if function_name:
+                return ct.models.MLModel(str(path), function_name=function_name)
+            else:
+                return ct.models.MLModel(str(path))
+                
+    except RuntimeError as e:
+        if "valid manifest does not exist" in str(e):
+            print(f"\nError: Could not load compiled model at {path}")
+            print("This might be because:")
+            print("1. The model is not properly compiled")
+            print("2. The model was compiled for a different OS version")
+            print("3. The model needs to be recompiled")
+            print("\nTry using the .mlpackage version instead, or recompile the model.")
+        raise
+
+def parse_args():
+    parser = argparse.ArgumentParser(description='Full Chat with CoreML LLaMA with context window shifting, gil resolved (c) 2025 Anemll')
+    
+    # Add meta.yaml option
+    parser.add_argument('--meta', type=str, help='Path to meta.yaml to load all parameters')
+    
+    # Add existing arguments
+    parser.add_argument('--d', '--dir', type=str, default='.',
+                       help='Directory containing model files (default: current directory)')
+    parser.add_argument('--embed', type=str, required=False,
+                       help='Path to embeddings model (relative to --dir)')
+    parser.add_argument('--ffn', type=str, required=False,
+                       help='Path to FFN model (can be chunked, relative to --dir)')
+    parser.add_argument('--lmhead', type=str, required=False,
+                       help='Path to LM head model (relative to --dir)')
+    parser.add_argument('--tokenizer', type=str, required=False,
+                       help='Path to tokenizer')
+    
+    # Add new argument for auto-generation
+    parser.add_argument('--prompt', type=str,
+                       help='If specified, run once with this prompt and exit')
+    
+    # Add no-warmup flag
+    parser.add_argument('--nw', action='store_true',
+                       help='Skip warmup phase')
+    
+    # Add debug level
+    parser.add_argument('--debug-level', type=int, default=0,
+                       help='Debug level (0=none, 1=print prompts, 2=more verbose)')
+    
+    # Model configuration
+    parser.add_argument('--context-length', type=int,
+                       help='Context length for the model (default: 512), if not provided, it will be detected from the model directory name ctxNUMBER')
+    parser.add_argument('--batch-size', type=int,
+                       help='Batch size for prefill (default: 64)')
+    
+    args = parser.parse_args()
+    
+    # If meta.yaml is provided, load parameters from it
+    if args.meta:
+        try:
+            with open(args.meta, 'r') as f:
+                meta = yaml.safe_load(f)
+            params = meta['model_info']['parameters']
+            
+            # Set model directory to meta.yaml directory if not specified
+            if not args.d or args.d == '.':
+                args.d = str(Path(args.meta).parent)
+            
+            # Build model paths based on parameters
+            prefix = params.get('model_prefix', 'llama')  # Default to 'llama' if not specified
+            lut_ffn = f"_lut{params['lut_ffn']}" if params['lut_ffn'] != 'none' else ''
+            lut_lmhead = f"_lut{params['lut_lmhead']}" if params['lut_lmhead'] != 'none' else ''
+            lut_embeddings = f"_lut{params['lut_embeddings']}" if params['lut_embeddings'] != 'none' else ''
+            num_chunks = int(params['num_chunks'])
+            
+            # Set model paths if not specified
+            if not args.lmhead:
+                args.lmhead = f'{prefix}_lm_head{lut_lmhead}'
+            if not args.embed:
+                args.embed = f'{prefix}_embeddings{lut_embeddings}'  # Changed from lm_head to embeddings
+            if not args.ffn:
+                args.ffn = f'{prefix}_FFN_PF{lut_ffn}_chunk_01of{num_chunks:02d}'
+            if not args.tokenizer:
+                args.tokenizer = args.d
+            
+            # Set other parameters if not overridden by command line
+            if args.context_length is None:
+                args.context_length = int(params['context_length'])
+            if args.batch_size is None:
+                args.batch_size = int(params['batch_size'])
+            args.num_chunks = num_chunks
+            
+            # Parse split_lm_head parameter from meta.yaml
+            if 'split_lm_head' in params:
+                args.split_lm_head = int(params['split_lm_head'])
+            else:
+                args.split_lm_head = 8  # Default value
+            
+            print(f"\nLoaded parameters from {args.meta}:")
+            print(f"  Context Length: {args.context_length}")
+            print(f"  Batch Size: {args.batch_size}")
+            print(f"  Num Chunks: {args.num_chunks}")
+            print(f"  Split LM Head: {args.split_lm_head}")
+            print(f"  Models Directory: {args.d}")
+            print(f"  Embeddings: {args.embed}")
+            print(f"  LM Head: {args.lmhead}")
+            print(f"  FFN: {args.ffn}")
+            
+        except Exception as e:
+            print(f"\nError loading meta.yaml: {str(e)}")
+            sys.exit(1)
+    
+    return args
+
+def load_metadata(model,args):
+    # Extract metadata and config parameters
+    metadata = {}
+    if hasattr(model, 'user_defined_metadata'):
+        meta = model.user_defined_metadata
+        
+        # Extract key parameters with defaults
+        metadata['context_length'] = int(meta.get('com.anemll.context_length', 512))
+        metadata['state_length'] = int(meta.get('com.anemll.state_length', metadata['context_length']))  # Added state_length
+        metadata['batch_size'] = int(meta.get('com.anemll.batch_size', 64))
+        metadata['lut_bits'] = int(meta.get('com.anemll.lut_bits', 0))
+        metadata['num_chunks'] = int(meta.get('com.anemll.num_chunks', 1))
+        
+        print("\nExtracted Parameters:")
+        print(f"  Context Length: {metadata['context_length']}")
+        print(f"  State Length: {metadata['state_length']}")
+        print(f"  Prefill Batch Size: {metadata['batch_size']}")
+        print(f"  LUT Bits: {metadata['lut_bits']}")
+        print(f"  Number of Chunks: {metadata['num_chunks']}")
+        
+        # Print model info
+        print("\nModel Info:")
+        if 'com.anemll.info' in meta:
+            print(f"  {meta['com.anemll.info']}")
+        if 'com.github.apple.coremltools.version' in meta:
+            print(f"  CoreML Tools: {meta['com.github.apple.coremltools.version']}")
+        
+        # Print model input/output shapes
+        print("\nModel Shapes:")
+        if hasattr(model, 'input_description'):
+            print("  Inputs:")
+            for name, desc in model.input_description.items():
+                print(f"    {name}: {desc}")
+        if hasattr(model, 'output_description'):
+            print("  Outputs:")
+            for name, desc in model.output_description.items():
+                print(f"    {name}: {desc}")
+    else:
+        print("\nWarning: No metadata found in model")
+
+        # Check if model directory name contains context length pattern (ctxXXX)
+        ctx_len = 512
+        if args.context_length is  None:
+            import re
+            ctx_match = re.search(r'ctx(\d+)', str(args.d))
+            if ctx_match:
+                ctx_len0 = int(ctx_match.group(1))
+                if 512 <= ctx_len0 <= 8096:
+                    ctx_len = ctx_len0
+                    print(f"\nDetected context length {ctx_len} from directory name")
+            else:
+                print(f"\nWarning: No context length found in directory  {ctx_len} from directory name {args.d}")
+        else:
+            ctx_len = args.context_length
+
+        # Use defaults or values from args
+        metadata['context_length'] = ctx_len
+        metadata['state_length'] = ctx_len
+        # Get batch size from args or use default
+        metadata['batch_size'] = getattr(args, 'batch_size', 64)
+        metadata['lut_bits'] = 4
+        metadata['num_chunks'] = getattr(args, 'num_chunks', 4)
+        print("\nUsing parameters:")
+        print(f"  Context Length: {metadata['context_length']}")
+        print(f"  State Length: {metadata['state_length']}")
+        print(f"  Prefill Batch Size: {metadata['batch_size']}")
+        print(f"  LUT Bits: {metadata['lut_bits']}")
+        print(f"  Number of Chunks: {metadata['num_chunks']}")
+
+    # Override with values from args if they exist
+    if hasattr(args, 'batch_size') and args.batch_size is not None:
+        metadata['batch_size'] = args.batch_size
+        print(f"\nOverriding batch size from args: {args.batch_size}")
+    if hasattr(args, 'num_chunks') and args.num_chunks is not None:
+        metadata['num_chunks'] = args.num_chunks
+        print(f"\nOverriding num chunks from args: {args.num_chunks}")
+    
+    return metadata
+    
+def load_models(args,metadata):
+    """Load all required models and extract metadata."""
+    print("\nLoading models...")
+    
+    try:
+        # Load embeddings model
+        print("\nLoading embeddings model...")
+        embed_path = parse_model_path(args.embed)
+        print(f"Loading from: {embed_path}")
+        embed_model = load_model(embed_path)
+        print("Embeddings model loaded successfully")
+        metadata = load_metadata(embed_model,args)
+        
+
+        
+        # Load LM head model
+        print("\nLoading LM head model...")
+        lmhead_path = parse_model_path(args.lmhead)
+        print(f"Loading from: {lmhead_path}")
+        lmhead_model = load_model(lmhead_path)
+        print("LM head model loaded successfully")
+        
+        # Parse FFN path and find chunks if needed
+        print("\nLoading FFN+PREFILL model(s)...")
+        ffn_path = parse_model_path(args.ffn)
+        chunk_no, total_chunks = parse_ffn_filename(ffn_path)
+        
+        ffn_models = []
+        if chunk_no and total_chunks:
+            print(f"\nDetected chunked FFN+PREFILL model ({total_chunks} chunks)")
+            # Find and load all chunks
+            chunk_paths = find_all_chunks(ffn_path)
+            if len(chunk_paths) != total_chunks:
+                raise ValueError(f"Found {len(chunk_paths)} chunks but filename indicates {total_chunks} chunks")
+                
+            for chunk_path in chunk_paths:
+                print(f"\nLoading FFN+PREFILL chunk: {Path(chunk_path).name}")
+                try:
+                    # For chunked models, we need both infer and prefill functions
+                    ffn_models.append({
+                        'infer': load_model(chunk_path, function_name='infer'),
+                        'prefill': load_model(chunk_path, function_name='prefill')
+                    })
+                    print("Chunk loaded successfully")
+                except Exception as e:
+                    print(f"Error loading chunk {chunk_path}: {str(e)}")
+                    raise
+            metadata = load_metadata(ffn_models[0],args)
+
+        else:
+            print("\nLoading single FFN model...")
+            ffn_models.append(load_model(ffn_path))
+            print("FFN model loaded successfully")
+        
+        return embed_model, ffn_models, lmhead_model, metadata
+        
+    except Exception as e:
+        print(f"\nError loading models: {str(e)}")
+        print("\nPlease ensure all model files exist and are accessible.")
+        print("Expected files:")
+        print(f"  Embeddings: {args.embed}")
+        print(f"  LM Head: {args.lmhead}")
+        print(f"  FFN: {args.ffn}")
+        raise
+
+# At the top of the file, make this a default path
+
+def initialize_tokenizer(model_path=None):
+    """Initialize and configure the tokenizer."""
+    try:
+
+        
+        tokenizer = AutoTokenizer.from_pretrained(
+            str(model_path), 
+            use_fast=False,
+            trust_remote_code=True
+        )
+        
+        print("\nTokenizer Configuration:")
+        print(f"Tokenizer type: {type(tokenizer)}")
+        print(f"Tokenizer name: {tokenizer.__class__.__name__}")
+        print(f"Vocabulary size: {len(tokenizer)}")
+        print(f"Model max length: {tokenizer.model_max_length}")
+
+        if tokenizer.pad_token is None:
+            tokenizer.pad_token = tokenizer.eos_token
+            tokenizer.pad_token_id = tokenizer.eos_token_id
+            print("Set PAD token to EOS token")
+        
+        tokenizer.padding_side = "left"
+        
+        print(f"\nSpecial Tokens:")
+        print(f"PAD token: '{tokenizer.pad_token}' (ID: {tokenizer.pad_token_id})")
+        print(f"EOS token: '{tokenizer.eos_token}' (ID: {tokenizer.eos_token_id})")
+        print(f"BOS token: '{tokenizer.bos_token}' (ID: {tokenizer.bos_token_id})")
+        print(f"UNK token: '{tokenizer.unk_token}' (ID: {tokenizer.unk_token_id})")
+
+        return tokenizer
+        
+    except Exception as e:
+        print(f"\nError: Failed to load tokenizer from {model_path}")
+        print(f"Error details: {str(e)}")
+        print(f"Error type: {type(e)}")
+        print("\nThis code requires a Llama 3.2 model for chat template functionality.")
+        print("Please provide the path to a Llama 3.2 model directory.")
+        import traceback
+        traceback.print_exc()
+        raise
+
+
+
+def make_causal_mask(length, start):
+    """Create causal attention mask."""
+    mask = np.full((1, 1, length, length), -np.inf, dtype=np.float16)
+    row_indices = np.arange(length).reshape(length, 1)
+    col_indices = np.arange(length).reshape(1, length)
+    mask[:, :, col_indices <= (row_indices + start)] = 0
+    return mask
+
+def run_prefill(embed_model, ffn_models, input_ids, current_pos, context_length, batch_size, state, causal_mask):
+    """Run prefill on the input sequence."""
+    #print(f"[DEBUG] Running prefill from 0 to {current_pos}")
+    
+    # Process in batches
+    batch_pos = 0
+    while batch_pos < current_pos:
+        batch_end = min(batch_pos + batch_size, current_pos)
+        current_batch_size = batch_end - batch_pos
+        
+        #print(f"[DEBUG] Prefill batch {batch_pos}-{batch_end} (size={current_batch_size})")
+        
+        # Get current batch
+        batch_input = input_ids[:, batch_pos:batch_end]
+        
+        # Pad to full batch size
+        batch_input = F.pad(
+            batch_input,
+            (0, batch_size - current_batch_size),
+            value=0
+        )
+        
+        # Generate position IDs for this batch
+        position_ids = torch.arange(batch_pos, batch_pos + batch_size, dtype=torch.int32)
+        
+        # Use the pre-initialized causal mask and extract the batch portion
+        batch_causal_mask = causal_mask[:, :, batch_pos:batch_pos + batch_size, :]
+        
+        # Run embeddings
+        hidden_states = torch.from_numpy(
+            embed_model.predict({'input_ids': batch_input.numpy()})['hidden_states']
+        )
+        
+        # Run through FFN chunks
+        for ffn_model in ffn_models:
+            if isinstance(ffn_model, dict):
+                inputs = {
+                    'hidden_states': hidden_states.numpy(),
+                    'position_ids': position_ids.numpy(),
+                    'causal_mask': batch_causal_mask.numpy(),
+                    'current_pos': np.array([batch_pos], dtype=np.int32)
+                }
+                output = ffn_model['prefill'].predict(inputs, state)
+                hidden_states = torch.from_numpy(output['output_hidden_states'])
+        
+        batch_pos = batch_end
+    
+    return torch.tensor([current_pos], dtype=torch.int32)
+
+def generate_next_token(embed_model, ffn_models, lmhead_model, input_ids, pos, context_length, state, causal_mask, metadata=None, temperature=0.0):
+    """Generate the next token."""
+    # Get current token
+    current_token = input_ids[:, pos-1:pos]
+    
+    # Run embeddings
+    hidden_states = torch.from_numpy(
+        embed_model.predict({'input_ids': current_token.numpy()})['hidden_states']
+    )
+    
+    # Create masks
+    update_mask = torch.zeros((1, 1, context_length, 1), dtype=torch.float16)
+    update_mask[0, 0, pos-1, 0] = 1.0
+    position_ids = torch.tensor([pos-1], dtype=torch.int32)
+    
+    # Use the pre-initialized causal mask and extract the single position portion
+    single_causal_mask = causal_mask[:, :, pos-1:pos, :]
+    
+    # Run through FFN chunks
+    for ffn_model in ffn_models:
+        if isinstance(ffn_model, dict):
+            inputs = {
+                'hidden_states': hidden_states.numpy(),
+                'update_mask': update_mask.numpy(),
+                'position_ids': position_ids.numpy(),
+                'causal_mask': single_causal_mask.numpy(),
+                'current_pos': position_ids.numpy()
+            }
+            output = ffn_model['infer'].predict(inputs, state)
+            hidden_states = torch.from_numpy(output['output_hidden_states'])
+    
+    # Run LM head and get next token
+    lm_output = lmhead_model.predict({'hidden_states': hidden_states.numpy()})
+    
+    if 'logits1' in lm_output:
+        logits_parts = []
+        for i in range(1, metadata.get('split_lm_head', 8) + 1):
+            key = f'logits{i}'
+            if key in lm_output:
+                logits_parts.append(torch.from_numpy(lm_output[key]))
+        logits = torch.cat(logits_parts, dim=-1)
+    else:
+        logits = torch.from_numpy(lm_output['output_logits'])
+    
+    if temperature > 0:
+        logits = logits / temperature
+        probs = F.softmax(logits[0, -1, :], dim=-1)
+        next_token = torch.multinomial(probs, num_samples=1).item()
+    else:
+        next_token = torch.argmax(logits[0, -1, :]).item()
+    
+    return next_token
+
+def create_unified_state(ffn_models, context_length):
+    """Create unified KV cache state for transformer."""
+    if isinstance(ffn_models[0], dict):
+        # Use first FFN model's prefill function to create state
+        state = ffn_models[0]['prefill'].make_state()
+        print(f"\nCreated unified transformer state for {len(ffn_models)} chunks")
+        return state
+    else:
+        state = ffn_models[0].make_state()
+        print("\nCreated unified transformer state")
+        return state
+
+def initialize_causal_mask(context_length):
+    """Initialize causal mask for transformer attention."""
+    causal_mask = make_causal_mask(context_length, 0)
+    causal_mask = torch.tensor(causal_mask, dtype=torch.float16)
+    print(f"\nInitialized causal mask for context length {context_length}")
+    return causal_mask
+
+def get_user_input():
+    """Get input from user, handling special key combinations."""
+    global THINKING_MODE
+    try:
+        import termios
+        import tty
+        import sys
+
+        def _getch():
+            fd = sys.stdin.fileno()
+            old_settings = termios.tcgetattr(fd)
+            try:
+                tty.setraw(sys.stdin.fileno())
+                ch = sys.stdin.read(1)
+            finally:
+                termios.tcsetattr(fd, termios.TCSADRAIN, old_settings)
+            return ch
+
+        buffer = []
+        while True:
+            char = _getch()
+            
+            # Debug: print the character code
+            print(f"\nKey pressed: {repr(char)} (hex: {hex(ord(char))})")
+            
+            # Check for Enter key
+            if char == '\r' or char == '\n':
+                print()  # Move to next line
+                input_text = ''.join(buffer)
+                # Check if the command is /t
+                if input_text == '/t':
+                    THINKING_MODE = not THINKING_MODE
+                    print(f"Thinking mode {'ON' if THINKING_MODE else 'OFF'}")
+                    buffer = []  # Clear buffer
+                    print(f"\n{LIGHT_GREEN}You{' (thinking)' if THINKING_MODE else ''}:{RESET_COLOR}", end=' ', flush=True)
+                    continue
+                return input_text
+                
+            # Handle backspace
+            if char == '\x7f':  # backspace
+                if buffer:
+                    buffer.pop()
+                    sys.stdout.write('\b \b')  # Erase character
+                    sys.stdout.flush()
+                continue
+                
+            # Handle Ctrl-C
+            if char == '\x03':  # Ctrl-C
+                print("^C")
+                raise KeyboardInterrupt
+                
+            # Print character and add to buffer
+            sys.stdout.write(char)
+            sys.stdout.flush()
+            buffer.append(char)
+            
+    except ImportError:
+        # Fallback for systems without termios
+        return input("> ")
+
+def chat_loop(embed_model, ffn_models, lmhead_model, tokenizer, metadata, state, causal_mask, auto_prompt=None, warmup=False):
+    """Interactive chat loop."""
+    global THINKING_MODE
+    global DEBUG_LEVEL
+    context_length = metadata.get('context_length')
+    batch_size = metadata.get('batch_size', 64)
+    
+    if not warmup:
+        print(f"\nUsing context length: {context_length}")
+        print("\nStarting chat session. Press Ctrl+D to exit.")
+        print("Type your message and press Enter to chat. Use /t to toggle thinking mode.")
+        print(f"Thinking mode is {'ON' if THINKING_MODE else 'OFF'}")
+    
+    # Keep track of conversation history
+    conversation = []
+    
+    try:
+        while True:
+            try:
+                if not warmup:
+                    print(f"\n{LIGHT_GREEN}You{' (thinking)' if THINKING_MODE else ''}:{RESET_COLOR}", end=' ', flush=True)
+                if auto_prompt is not None:
+                    user_input = auto_prompt
+                    if not warmup:
+                        print(user_input)
+                else:
+                    user_input = input().strip()
+            except EOFError:
+                if not warmup:
+                    print("\nExiting chat...")
+                break
+            
+            if not user_input:
+                continue
+
+            # Handle /t command
+            if user_input == "/t":
+                THINKING_MODE = not THINKING_MODE
+                print(f"Thinking mode {'ON' if THINKING_MODE else 'OFF'}")
+                continue
+            
+            # Add user message to conversation
+            conversation.append({"role": "user", "content": user_input})
+            
+            # Format using chat template with full history
+            if THINKING_MODE:
+                # Add thinking prompt to system message
+                conversation_with_thinking = [{"role": "system", "content": THINKING_PROMPT}] + conversation
+                base_input_ids = tokenizer.apply_chat_template(
+                    conversation_with_thinking,
+                    return_tensors="pt",
+                    add_generation_prompt=True
+                ).to(torch.int32)
+                
+                # Print full prompt if debug level >= 1
+                if DEBUG_LEVEL >= 1 and not warmup:
+                    print(f"\n{DARK_BLUE}Debug: Full prompt with thinking:{RESET_COLOR}")
+                    print(tokenizer.decode(base_input_ids[0]))
+            else:
+                base_input_ids = tokenizer.apply_chat_template(
+                    conversation,
+                    return_tensors="pt",
+                    add_generation_prompt=True
+                ).to(torch.int32)
+                
+                # Print full prompt if debug level >= 1
+                if DEBUG_LEVEL >= 1 and not warmup:
+                    print(f"\n{DARK_BLUE}Debug: Full prompt:{RESET_COLOR}")
+                    print(tokenizer.decode(base_input_ids[0]))
+            
+            # Check if we need to trim history
+            while base_input_ids.size(1) > context_length - 100:  # Leave room for response
+                # Remove oldest message pair (user + assistant)
+                if len(conversation) > 2:
+                    conversation = conversation[2:]  # Remove oldest pair
+                    base_input_ids = tokenizer.apply_chat_template(
+                        conversation,
+                        return_tensors="pt",
+                        add_generation_prompt=True
+                    ).to(torch.int32)
+                else:
+                    # If only current message remains and still too long, truncate
+                    base_input_ids = base_input_ids[:, -context_length//2:]
+                    break
+            
+            context_pos = base_input_ids.size(1)
+            
+            # Pad sequence to context_size
+            input_ids = F.pad(
+                base_input_ids,
+                (0, context_length - context_pos),
+                value=0
+            )
+            
+            if not warmup:
+                print(f"\n{LIGHT_BLUE}Assistant:{RESET_COLOR}", end=' ', flush=True)
+            
+            # split_lm_head should already be in metadata from caller
+            
+            # Initialize token printer and collect response
+            token_printer = TokenPrinter(tokenizer)
+            response_tokens = []
+            generation_start_time = time.time()
+            
+            try:
+                # Run prefill on entire context
+                current_pos = run_prefill(
+                    embed_model,
+                    ffn_models,
+                    input_ids,
+                    context_pos,
+                    context_length,
+                    batch_size,
+                    state,
+                    causal_mask
+                )
+                #print(f"\n[DEBUG] After initial prefill - current_pos: {current_pos}")
+                
+                # Generation loop
+                pos = context_pos
+                tokens_generated = 0
+                inference_start = time.time()  # Start inference timing
+                
+                while True:
+                    # Check if we need to shift window
+                    if pos >= context_length - 2:
+                        # Calculate shift to maintain full batches
+                        batch_size = metadata.get('batch_size', 64)
+                        # Calculate max batches that fit in context
+                        max_batches = context_length // batch_size
+                        desired_batches = max(1, max_batches - 2)  # Leave room for new tokens
+                        new_size = min(desired_batches * batch_size, context_length - batch_size)
+                        
+                        # Create shifted input_ids
+                        tmp = torch.zeros((1, context_length), dtype=torch.int32)
+                        tmp[:,0:new_size] = input_ids[:,pos-new_size:pos]
+                        input_ids = tmp
+                        
+                        # Reset state and run prefill
+                        # keep the same state
+                        #state = create_unified_state(ffn_models, context_length)
+                        current_pos = run_prefill(
+                            embed_model,
+                            ffn_models,
+                            input_ids,
+                            new_size,  # Prefill the entire shifted content
+                            context_length,
+                            batch_size,
+                            state,
+                            causal_mask
+                        )
+                        
+                        # Start generating from the next position
+                        pos = new_size  # Don't back up, continue from where we left off
+                        
+                        #print(f"\n[DEBUG] After shift - next token will be at pos {pos}")
+                        #print(f"[DEBUG] Context before next token: {tokenizer.decode(input_ids[0, pos-40:pos])}")
+                        
+                        window_shifted = True
+                    
+                    # Generate next token
+                    next_token = generate_next_token(
+                        embed_model,
+                        ffn_models,
+                        lmhead_model,
+                        input_ids,
+                        pos,
+                        context_length,
+                        state,
+                        causal_mask,
+                        metadata
+                    )
+                    
+                    # Add token
+                    input_ids[0, pos] = next_token
+                    if not warmup:
+                        token_printer.add_token(next_token)
+                        token_printer.drain_buffer()
+                    response_tokens.append(next_token)
+                    
+                    pos += 1
+                    tokens_generated += 1
+                    
+                    # In warmup mode, limit tokens
+                    if warmup and tokens_generated >= WARMUP_TOKEN_LIMIT:
+                        break
+                    
+                    if next_token == tokenizer.eos_token_id:
+                        break
+                
+                inference_time = time.time() - inference_start  # Calculate inference time
+                
+                # Add assistant response to conversation
+                response_text = token_printer.stop()
+                conversation.append({"role": "assistant", "content": response_text})
+                
+                # Print stats only if not in warmup
+                if not warmup:
+                    total_time = time.time() - generation_start_time
+                    prefill_time = total_time - inference_time
+                    inference_tokens_per_sec = len(response_tokens) / inference_time if inference_time > 0 else 0
+                    prefill_ms = prefill_time * 1000
+                    prefill_tokens_per_sec = context_pos / prefill_time if prefill_time > 0 else 0
+                    print(f"{DARK_BLUE}{inference_tokens_per_sec:.1f} t/s, "
+                          f"TTFT: {prefill_ms:.1f}ms ({prefill_tokens_per_sec:.1f} t/s), "
+                          f"{len(response_tokens)} tokens{RESET_COLOR}")
+                
+                if auto_prompt is not None:
+                    break
+                
+            except KeyboardInterrupt:
+                if not warmup:
+                    print("\nGeneration interrupted")
+                token_printer.stop()
+                continue
+                
+    except Exception as e:
+        if not warmup:
+            print(f"\nError in chat loop: {str(e)}")
+            import traceback
+            traceback.print_exc()
+
+def main():
+    args = parse_args()
+    global DEBUG_LEVEL
+    DEBUG_LEVEL = args.debug_level
+    
+    # Convert directory to absolute path
+    model_dir = Path(args.d).resolve()
+    if not model_dir.exists():
+        print(f"\nError: Model directory not found: {model_dir}")
+        return 1
+        
+    print(f"\nUsing model directory: {model_dir}")
+    print(f"Context length: {args.context_length}")
+    
+    try:
+        # Update paths to be relative to model directory
+        args.embed = str(model_dir / args.embed)
+        args.ffn = str(model_dir / args.ffn)
+        args.lmhead = str(model_dir / args.lmhead)
+        
+        # Handle tokenizer path separately since it's not relative to model_dir
+        if args.tokenizer is None:
+            args.tokenizer = str(model_dir)
+        
+        if not Path(args.tokenizer).exists():
+            print(f"\nError: Tokenizer directory not found: {args.tokenizer}")
+            return 1
+    
+        args.tokenizer = str(Path(args.tokenizer).resolve())  # Convert to absolute path
+        print(f"Using tokenizer path: {args.tokenizer}")
+        
+        metadata = {}
+        # Load models and extract metadata
+        embed_model, ffn_models, lmhead_model, metadata = load_models(args,metadata)
+        
+        print(f"\nMetadata befor args.context_length: {metadata}")
+
+        # Override context length from command line if provided
+        if args.context_length is not None:
+            metadata['context_length'] = args.context_length
+            metadata['state_length'] = args.context_length  # Also update state_length
+            print(f"\nOverriding context length from command line: {args.context_length}")
+        
+        print(f"\nMetadata after load_models: {metadata}")
+        
+        # Load tokenizer with resolved path
+        tokenizer = initialize_tokenizer(args.tokenizer)
+        if tokenizer is None:
+            raise RuntimeError("Failed to initialize tokenizer")
+        
+        # Create unified state once
+        state = create_unified_state(ffn_models, metadata['context_length'])
+        
+        # Initialize causal mask once
+        causal_mask = initialize_causal_mask(metadata['context_length'])
+        
+        # Add split_lm_head to metadata for generate_next_token
+        metadata['split_lm_head'] = getattr(args, 'split_lm_head', 8)
+        
+        # Warmup runs to prevent Python GIL issues with CoreML !
+        if not args.nw:
+            for i in range(2):
+                chat_loop(
+                    embed_model=embed_model,
+                    ffn_models=ffn_models,
+                    lmhead_model=lmhead_model,
+                    tokenizer=tokenizer,
+                    metadata=metadata,
+                    state=state,  # Pass the state
+                    causal_mask=causal_mask,  # Pass the causal mask
+                    warmup=True,
+                    auto_prompt="who are you?"
+                )
+        
+        # Main run
+        chat_loop(
+            embed_model=embed_model,
+            ffn_models=ffn_models,
+            lmhead_model=lmhead_model,
+            tokenizer=tokenizer,
+            metadata=metadata,
+            state=state,  # Pass the state
+            causal_mask=causal_mask,  # Pass the causal mask
+            warmup=False,
+            auto_prompt=args.prompt
+        )
+        
+    except Exception as e:
+        print(f"\nError: {str(e)}")
+        import traceback
+        traceback.print_exc()
+        return 1
+    
+    return 0
+
+if __name__ == "__main__":
+    exit(main()) 

config.json

@@ -0,0 +1,4 @@
+{
+  "tokenizer_class": "LlamaTokenizer",
+  "model_type": "llama"
+}

meta.yaml

@@ -0,0 +1,24 @@
+model_info:
+  name: anemll-Qwen3-4B-MLX-dequantized-ctx1024
+  version: 0.3.0
+  description: |
+    Demonstarates running Qwen3-4B-MLX-dequantized on Apple Neural Engine
+    Context length: 1024
+    Batch size: 64
+    Chunks: 2
+  license: MIT
+  author: Anemll
+  framework: Core ML
+  language: Python
+  parameters:
+    context_length: 1024
+    batch_size: 64
+    lut_embeddings: none
+    lut_ffn: 4
+    lut_lmhead: 8
+    num_chunks: 2
+    model_prefix: qwen
+    embeddings: qwen_embeddings.mlmodelc
+    lm_head: qwen_lm_head_lut8.mlmodelc
+    ffn: qwen_FFN_PF_lut4.mlmodelc
+    split_lm_head: 16

qwen_FFN_PF_lut4_chunk_01of02.mlmodelc/analytics/coremldata.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f940f9661b267bb9282e47cd2ba5c999e4a76de57404fcc5a9f4a1b964db4863
+size 243

qwen_FFN_PF_lut4_chunk_01of02.mlmodelc/coremldata.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a7bea695683f198193ea833dd4dc7559ac47137a672afd457a5234edac4545da
+size 983

qwen_FFN_PF_lut4_chunk_01of02.mlmodelc/metadata.json

@@ -0,0 +1,324 @@
+[
+  {
+    "metadataOutputVersion" : "3.0",
+    "userDefinedMetadata" : {
+      "com.github.apple.coremltools.source" : "torch==2.5.0",
+      "com.github.apple.coremltools.version" : "8.3.0",
+      "com.anemll.context_length" : "1024",
+      "com.anemll.chunk_no" : "1",
+      "com.github.apple.coremltools.source_dialect" : "TorchScript",
+      "com.anemll.num_chunks" : "2",
+      "com.anemll.info" : "Converted with Anemll v0.3.0",
+      "com.anemll.batch_size" : "64",
+      "com.anemll.lut_bits" : "4"
+    },
+    "availability" : {
+      "macOS" : "15.0",
+      "tvOS" : "18.0",
+      "visionOS" : "2.0",
+      "watchOS" : "11.0",
+      "iOS" : "18.0",
+      "macCatalyst" : "18.0"
+    },
+    "inputSchema" : [
+      {
+        "hasShapeFlexibility" : "0",
+        "isOptional" : "0",
+        "dataType" : "Float16",
+        "formattedType" : "MultiArray (Float16 1 × 1 × 2560)",
+        "shortDescription" : "",
+        "shape" : "[1, 1, 2560]",
+        "name" : "hidden_states",
+        "type" : "MultiArray"
+      },
+      {
+        "hasShapeFlexibility" : "0",
+        "isOptional" : "0",
+        "dataType" : "Int32",
+        "formattedType" : "MultiArray (Int32 1)",
+        "shortDescription" : "",
+        "shape" : "[1]",
+        "name" : "position_ids",
+        "type" : "MultiArray"
+      },
+      {
+        "hasShapeFlexibility" : "0",
+        "isOptional" : "0",
+        "dataType" : "Float16",
+        "formattedType" : "MultiArray (Float16 1 × 1 × 1 × 1024)",
+        "shortDescription" : "",
+        "shape" : "[1, 1, 1, 1024]",
+        "name" : "causal_mask",
+        "type" : "MultiArray"
+      },
+      {
+        "hasShapeFlexibility" : "0",
+        "isOptional" : "0",
+        "dataType" : "Int32",
+        "formattedType" : "MultiArray (Int32 1)",
+        "shortDescription" : "",
+        "shape" : "[1]",
+        "name" : "current_pos",
+        "type" : "MultiArray"
+      }
+    ],
+    "outputSchema" : [
+      {
+        "hasShapeFlexibility" : "0",
+        "isOptional" : "0",
+        "dataType" : "Float16",
+        "formattedType" : "MultiArray (Float16 1 × 1 × 2560)",
+        "shortDescription" : "",
+        "shape" : "[1, 1, 2560]",
+        "name" : "output_hidden_states",
+        "type" : "MultiArray"
+      }
+    ],
+    "modelParameters" : [
+
+    ],
+    "storagePrecision" : "Mixed (Float16, Palettized (11 bits), Palettized (13 bits), Palettized (15 bits), UInt4)",
+    "method" : "predict",
+    "functions" : [
+      {
+        "inputSchema" : [
+          {
+            "hasShapeFlexibility" : "0",
+            "isOptional" : "0",
+            "dataType" : "Float16",
+            "formattedType" : "MultiArray (Float16 1 × 1 × 2560)",
+            "shortDescription" : "",
+            "shape" : "[1, 1, 2560]",
+            "name" : "hidden_states",
+            "type" : "MultiArray"
+          },
+          {
+            "hasShapeFlexibility" : "0",
+            "isOptional" : "0",
+            "dataType" : "Int32",
+            "formattedType" : "MultiArray (Int32 1)",
+            "shortDescription" : "",
+            "shape" : "[1]",
+            "name" : "position_ids",
+            "type" : "MultiArray"
+          },
+          {
+            "hasShapeFlexibility" : "0",
+            "isOptional" : "0",
+            "dataType" : "Float16",
+            "formattedType" : "MultiArray (Float16 1 × 1 × 1 × 1024)",
+            "shortDescription" : "",
+            "shape" : "[1, 1, 1, 1024]",
+            "name" : "causal_mask",
+            "type" : "MultiArray"
+          },
+          {
+            "hasShapeFlexibility" : "0",
+            "isOptional" : "0",
+            "dataType" : "Int32",
+            "formattedType" : "MultiArray (Int32 1)",
+            "shortDescription" : "",
+            "shape" : "[1]",
+            "name" : "current_pos",
+            "type" : "MultiArray"
+          }
+        ],
+        "computePrecision" : "Mixed (Float16, Int32)",
+        "storagePrecision" : "Mixed (Float16, Palettized (11 bits), Palettized (13 bits), Palettized (15 bits), UInt4)",
+        "stateSchema" : [
+          {
+            "dataType" : "Float16",
+            "isOptional" : "0",
+            "formattedType" : "State (Float16 72 × 8 × 1024 × 128)",
+            "shortDescription" : "",
+            "shape" : "[72, 8, 1024, 128]",
+            "name" : "model_model_kv_cache_0",
+            "type" : "State"
+          }
+        ],
+        "outputSchema" : [
+          {
+            "hasShapeFlexibility" : "0",
+            "isOptional" : "0",
+            "dataType" : "Float16",
+            "formattedType" : "MultiArray (Float16 1 × 1 × 2560)",
+            "shortDescription" : "",
+            "shape" : "[1, 1, 2560]",
+            "name" : "output_hidden_states",
+            "type" : "MultiArray"
+          }
+        ],
+        "name" : "infer",
+        "mlProgramOperationTypeHistogram" : {
+          "Ios18.expandDims" : 72,
+          "Ios18.mul" : 144,
+          "Ios18.softmax" : 18,
+          "Ios18.matmul" : 36,
+          "Identity" : 1,
+          "Ios16.reduceMean" : 73,
+          "Ios18.greaterEqual" : 1,
+          "Select" : 1,
+          "Ios18.readState" : 37,
+          "Tile" : 36,
+          "Ios18.gather" : 2,
+          "Ios18.add" : 92,
+          "Ios18.layerNorm" : 73,
+          "Ios18.sliceUpdate" : 36,
+          "Ios18.writeState" : 36,
+          "Ios18.reshape" : 110,
+          "Ios18.constexprLutToDense" : 126,
+          "Ios18.conv" : 126,
+          "Ios18.concat" : 108,
+          "Ios18.transpose" : 108,
+          "Ios18.sub" : 73,
+          "Ios18.silu" : 18,
+          "Ios18.sliceByIndex" : 108,
+          "Ios18.squeeze" : 54
+        }
+      },
+      {
+        "inputSchema" : [
+          {
+            "hasShapeFlexibility" : "0",
+            "isOptional" : "0",
+            "dataType" : "Float16",
+            "formattedType" : "MultiArray (Float16 1 × 64 × 2560)",
+            "shortDescription" : "",
+            "shape" : "[1, 64, 2560]",
+            "name" : "hidden_states",
+            "type" : "MultiArray"
+          },
+          {
+            "hasShapeFlexibility" : "0",
+            "isOptional" : "0",
+            "dataType" : "Int32",
+            "formattedType" : "MultiArray (Int32 64)",
+            "shortDescription" : "",
+            "shape" : "[64]",
+            "name" : "position_ids",
+            "type" : "MultiArray"
+          },
+          {
+            "hasShapeFlexibility" : "0",
+            "isOptional" : "0",
+            "dataType" : "Float16",
+            "formattedType" : "MultiArray (Float16 1 × 1 × 64 × 1024)",
+            "shortDescription" : "",
+            "shape" : "[1, 1, 64, 1024]",
+            "name" : "causal_mask",
+            "type" : "MultiArray"
+          },
+          {
+            "hasShapeFlexibility" : "0",
+            "isOptional" : "0",
+            "dataType" : "Int32",
+            "formattedType" : "MultiArray (Int32 1)",
+            "shortDescription" : "",
+            "shape" : "[1]",
+            "name" : "current_pos",
+            "type" : "MultiArray"
+          }
+        ],
+        "computePrecision" : "Mixed (Float16, Int32)",
+        "storagePrecision" : "Mixed (Float16, Palettized (11 bits), Palettized (13 bits), Palettized (15 bits), UInt4)",
+        "stateSchema" : [
+          {
+            "dataType" : "Float16",
+            "isOptional" : "0",
+            "formattedType" : "State (Float16 72 × 8 × 1024 × 128)",
+            "shortDescription" : "",
+            "shape" : "[72, 8, 1024, 128]",
+            "name" : "model_model_kv_cache_0",
+            "type" : "State"
+          }
+        ],
+        "outputSchema" : [
+          {
+            "hasShapeFlexibility" : "0",
+            "isOptional" : "0",
+            "dataType" : "Float16",
+            "formattedType" : "MultiArray (Float16 1 × 64 × 2560)",
+            "shortDescription" : "",
+            "shape" : "[1, 64, 2560]",
+            "name" : "output_hidden_states",
+            "type" : "MultiArray"
+          }
+        ],
+        "name" : "prefill",
+        "mlProgramOperationTypeHistogram" : {
+          "Ios18.expandDims" : 72,
+          "Ios18.mul" : 144,
+          "Ios18.softmax" : 18,
+          "Ios18.matmul" : 36,
+          "Ios16.reduceMean" : 72,
+          "Ios18.greaterEqual" : 1,
+          "Select" : 1,
+          "Ios18.readState" : 37,
+          "Tile" : 36,
+          "Ios18.gather" : 2,
+          "Ios18.add" : 92,
+          "Ios18.layerNorm" : 72,
+          "Ios18.sliceUpdate" : 36,
+          "Ios18.writeState" : 36,
+          "Ios18.reshape" : 146,
+          "Ios18.constexprLutToDense" : 126,
+          "Ios18.conv" : 126,
+          "Ios18.concat" : 108,
+          "Ios18.transpose" : 164,
+          "Ios18.sub" : 72,
+          "Ios18.silu" : 18,
+          "Ios18.sliceByIndex" : 108,
+          "Ios18.squeeze" : 54
+        }
+      }
+    ],
+    "version" : "0.3.0",
+    "isUpdatable" : "0",
+    "defaultFunctionName" : "infer",
+    "specificationVersion" : 9,
+    "stateSchema" : [
+      {
+        "dataType" : "Float16",
+        "isOptional" : "0",
+        "formattedType" : "State (Float16 72 × 8 × 1024 × 128)",
+        "shortDescription" : "",
+        "shape" : "[72, 8, 1024, 128]",
+        "name" : "model_model_kv_cache_0",
+        "type" : "State"
+      }
+    ],
+    "computePrecision" : "Mixed (Float16, Int32)",
+    "mlProgramOperationTypeHistogram" : {
+      "Ios18.expandDims" : 72,
+      "Ios18.mul" : 144,
+      "Ios18.softmax" : 18,
+      "Ios18.matmul" : 36,
+      "Identity" : 1,
+      "Ios16.reduceMean" : 73,
+      "Ios18.greaterEqual" : 1,
+      "Select" : 1,
+      "Ios18.readState" : 37,
+      "Tile" : 36,
+      "Ios18.gather" : 2,
+      "Ios18.add" : 92,
+      "Ios18.layerNorm" : 73,
+      "Ios18.sliceUpdate" : 36,
+      "Ios18.writeState" : 36,
+      "Ios18.reshape" : 110,
+      "Ios18.constexprLutToDense" : 126,
+      "Ios18.conv" : 126,
+      "Ios18.concat" : 108,
+      "Ios18.transpose" : 108,
+      "Ios18.sub" : 73,
+      "Ios18.silu" : 18,
+      "Ios18.sliceByIndex" : 108,
+      "Ios18.squeeze" : 54
+    },
+    "shortDescription" : "Anemll Model: Multifunction FFN+Prefill",
+    "generatedClassName" : "qwen_FFN_PF_lut4_chunk_01of02",
+    "author" : "Converted with Anemll v0.3.0",
+    "modelType" : {
+      "name" : "MLModelType_mlProgram"
+    }
+  }
+]

qwen_FFN_PF_lut4_chunk_01of02.mlmodelc/weights/weight.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5f65eef202f5908da7c989174df68f29344a91520d01df382e89930fc8ac8140
+size 944314880

qwen_FFN_PF_lut4_chunk_02of02.mlmodelc/analytics/coremldata.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9b3e5221beac8057139104891eb8f4b3a2b53ad0a702aff31dc9b60fc0bd66c3
+size 243

qwen_FFN_PF_lut4_chunk_02of02.mlmodelc/coremldata.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:47a7d13c9ac8bd8b36abf6202b6b0f879030938a7af0b7022eb13f303f3ec5e6
+size 983

qwen_FFN_PF_lut4_chunk_02of02.mlmodelc/metadata.json

@@ -0,0 +1,324 @@
+[
+  {
+    "metadataOutputVersion" : "3.0",
+    "userDefinedMetadata" : {
+      "com.github.apple.coremltools.source" : "torch==2.5.0",
+      "com.github.apple.coremltools.source_dialect" : "TorchScript",
+      "com.github.apple.coremltools.version" : "8.3.0",
+      "com.anemll.chunk_no" : "2",
+      "com.anemll.context_length" : "1024",
+      "com.anemll.num_chunks" : "2",
+      "com.anemll.batch_size" : "64",
+      "com.anemll.info" : "Converted with Anemll v0.3.0",
+      "com.anemll.lut_bits" : "4"
+    },
+    "availability" : {
+      "macOS" : "15.0",
+      "tvOS" : "18.0",
+      "visionOS" : "2.0",
+      "watchOS" : "11.0",
+      "iOS" : "18.0",
+      "macCatalyst" : "18.0"
+    },
+    "inputSchema" : [
+      {
+        "hasShapeFlexibility" : "0",
+        "isOptional" : "0",
+        "dataType" : "Float16",
+        "formattedType" : "MultiArray (Float16 1 × 1 × 2560)",
+        "shortDescription" : "",
+        "shape" : "[1, 1, 2560]",
+        "name" : "hidden_states",
+        "type" : "MultiArray"
+      },
+      {
+        "hasShapeFlexibility" : "0",
+        "isOptional" : "0",
+        "dataType" : "Int32",
+        "formattedType" : "MultiArray (Int32 1)",
+        "shortDescription" : "",
+        "shape" : "[1]",
+        "name" : "position_ids",
+        "type" : "MultiArray"
+      },
+      {
+        "hasShapeFlexibility" : "0",
+        "isOptional" : "0",
+        "dataType" : "Float16",
+        "formattedType" : "MultiArray (Float16 1 × 1 × 1 × 1024)",
+        "shortDescription" : "",
+        "shape" : "[1, 1, 1, 1024]",
+        "name" : "causal_mask",
+        "type" : "MultiArray"
+      },
+      {
+        "hasShapeFlexibility" : "0",
+        "isOptional" : "0",
+        "dataType" : "Int32",
+        "formattedType" : "MultiArray (Int32 1)",
+        "shortDescription" : "",
+        "shape" : "[1]",
+        "name" : "current_pos",
+        "type" : "MultiArray"
+      }
+    ],
+    "outputSchema" : [
+      {
+        "hasShapeFlexibility" : "0",
+        "isOptional" : "0",
+        "dataType" : "Float16",
+        "formattedType" : "MultiArray (Float16 1 × 1 × 2560)",
+        "shortDescription" : "",
+        "shape" : "[1, 1, 2560]",
+        "name" : "output_hidden_states",
+        "type" : "MultiArray"
+      }
+    ],
+    "modelParameters" : [
+
+    ],
+    "storagePrecision" : "Mixed (Float16, Palettized (11 bits), Palettized (13 bits), Palettized (15 bits), UInt4)",
+    "method" : "predict",
+    "functions" : [
+      {
+        "inputSchema" : [
+          {
+            "hasShapeFlexibility" : "0",
+            "isOptional" : "0",
+            "dataType" : "Float16",
+            "formattedType" : "MultiArray (Float16 1 × 1 × 2560)",
+            "shortDescription" : "",
+            "shape" : "[1, 1, 2560]",
+            "name" : "hidden_states",
+            "type" : "MultiArray"
+          },
+          {
+            "hasShapeFlexibility" : "0",
+            "isOptional" : "0",
+            "dataType" : "Int32",
+            "formattedType" : "MultiArray (Int32 1)",
+            "shortDescription" : "",
+            "shape" : "[1]",
+            "name" : "position_ids",
+            "type" : "MultiArray"
+          },
+          {
+            "hasShapeFlexibility" : "0",
+            "isOptional" : "0",
+            "dataType" : "Float16",
+            "formattedType" : "MultiArray (Float16 1 × 1 × 1 × 1024)",
+            "shortDescription" : "",
+            "shape" : "[1, 1, 1, 1024]",
+            "name" : "causal_mask",
+            "type" : "MultiArray"
+          },
+          {
+            "hasShapeFlexibility" : "0",
+            "isOptional" : "0",
+            "dataType" : "Int32",
+            "formattedType" : "MultiArray (Int32 1)",
+            "shortDescription" : "",
+            "shape" : "[1]",
+            "name" : "current_pos",
+            "type" : "MultiArray"
+          }
+        ],
+        "computePrecision" : "Mixed (Float16, Int32)",
+        "storagePrecision" : "Mixed (Float16, Palettized (11 bits), Palettized (13 bits), Palettized (15 bits), UInt4)",
+        "stateSchema" : [
+          {
+            "dataType" : "Float16",
+            "isOptional" : "0",
+            "formattedType" : "State (Float16 72 × 8 × 1024 × 128)",
+            "shortDescription" : "",
+            "shape" : "[72, 8, 1024, 128]",
+            "name" : "model_model_kv_cache_0",
+            "type" : "State"
+          }
+        ],
+        "outputSchema" : [
+          {
+            "hasShapeFlexibility" : "0",
+            "isOptional" : "0",
+            "dataType" : "Float16",
+            "formattedType" : "MultiArray (Float16 1 × 1 × 2560)",
+            "shortDescription" : "",
+            "shape" : "[1, 1, 2560]",
+            "name" : "output_hidden_states",
+            "type" : "MultiArray"
+          }
+        ],
+        "name" : "infer",
+        "mlProgramOperationTypeHistogram" : {
+          "Ios18.expandDims" : 72,
+          "Ios18.mul" : 144,
+          "Ios18.softmax" : 18,
+          "Ios18.matmul" : 36,
+          "Identity" : 1,
+          "Ios16.reduceMean" : 73,
+          "Ios18.greaterEqual" : 1,
+          "Select" : 1,
+          "Ios18.readState" : 37,
+          "Tile" : 36,
+          "Ios18.gather" : 2,
+          "Ios18.add" : 92,
+          "Ios18.layerNorm" : 73,
+          "Ios18.sliceUpdate" : 36,
+          "Ios18.writeState" : 36,
+          "Ios18.reshape" : 110,
+          "Ios18.constexprLutToDense" : 126,
+          "Ios18.conv" : 126,
+          "Ios18.concat" : 108,
+          "Ios18.transpose" : 108,
+          "Ios18.sub" : 73,
+          "Ios18.silu" : 18,
+          "Ios18.sliceByIndex" : 108,
+          "Ios18.squeeze" : 54
+        }
+      },
+      {
+        "inputSchema" : [
+          {
+            "hasShapeFlexibility" : "0",
+            "isOptional" : "0",
+            "dataType" : "Float16",
+            "formattedType" : "MultiArray (Float16 1 × 64 × 2560)",
+            "shortDescription" : "",
+            "shape" : "[1, 64, 2560]",
+            "name" : "hidden_states",
+            "type" : "MultiArray"
+          },
+          {
+            "hasShapeFlexibility" : "0",
+            "isOptional" : "0",
+            "dataType" : "Int32",
+            "formattedType" : "MultiArray (Int32 64)",
+            "shortDescription" : "",
+            "shape" : "[64]",
+            "name" : "position_ids",
+            "type" : "MultiArray"
+          },
+          {
+            "hasShapeFlexibility" : "0",
+            "isOptional" : "0",
+            "dataType" : "Float16",
+            "formattedType" : "MultiArray (Float16 1 × 1 × 64 × 1024)",
+            "shortDescription" : "",
+            "shape" : "[1, 1, 64, 1024]",
+            "name" : "causal_mask",
+            "type" : "MultiArray"
+          },
+          {
+            "hasShapeFlexibility" : "0",
+            "isOptional" : "0",
+            "dataType" : "Int32",
+            "formattedType" : "MultiArray (Int32 1)",
+            "shortDescription" : "",
+            "shape" : "[1]",
+            "name" : "current_pos",
+            "type" : "MultiArray"
+          }
+        ],
+        "computePrecision" : "Mixed (Float16, Int32)",
+        "storagePrecision" : "Mixed (Float16, Palettized (11 bits), Palettized (13 bits), Palettized (15 bits), UInt4)",
+        "stateSchema" : [
+          {
+            "dataType" : "Float16",
+            "isOptional" : "0",
+            "formattedType" : "State (Float16 72 × 8 × 1024 × 128)",
+            "shortDescription" : "",
+            "shape" : "[72, 8, 1024, 128]",
+            "name" : "model_model_kv_cache_0",
+            "type" : "State"
+          }
+        ],
+        "outputSchema" : [
+          {
+            "hasShapeFlexibility" : "0",
+            "isOptional" : "0",
+            "dataType" : "Float16",
+            "formattedType" : "MultiArray (Float16 1 × 64 × 2560)",
+            "shortDescription" : "",
+            "shape" : "[1, 64, 2560]",
+            "name" : "output_hidden_states",
+            "type" : "MultiArray"
+          }
+        ],
+        "name" : "prefill",
+        "mlProgramOperationTypeHistogram" : {
+          "Ios18.expandDims" : 72,
+          "Ios18.mul" : 144,
+          "Ios18.softmax" : 18,
+          "Ios18.matmul" : 36,
+          "Ios16.reduceMean" : 73,
+          "Ios18.greaterEqual" : 1,
+          "Select" : 1,
+          "Ios18.readState" : 37,
+          "Tile" : 36,
+          "Ios18.gather" : 2,
+          "Ios18.add" : 92,
+          "Ios18.layerNorm" : 73,
+          "Ios18.sliceUpdate" : 36,
+          "Ios18.writeState" : 36,
+          "Ios18.reshape" : 146,
+          "Ios18.constexprLutToDense" : 126,
+          "Ios18.conv" : 126,
+          "Ios18.concat" : 108,
+          "Ios18.transpose" : 164,
+          "Ios18.sub" : 73,
+          "Ios18.silu" : 18,
+          "Ios18.sliceByIndex" : 108,
+          "Ios18.squeeze" : 54
+        }
+      }
+    ],
+    "version" : "0.3.0",
+    "isUpdatable" : "0",
+    "defaultFunctionName" : "infer",
+    "specificationVersion" : 9,
+    "stateSchema" : [
+      {
+        "dataType" : "Float16",
+        "isOptional" : "0",
+        "formattedType" : "State (Float16 72 × 8 × 1024 × 128)",
+        "shortDescription" : "",
+        "shape" : "[72, 8, 1024, 128]",
+        "name" : "model_model_kv_cache_0",
+        "type" : "State"
+      }
+    ],
+    "computePrecision" : "Mixed (Float16, Int32)",
+    "mlProgramOperationTypeHistogram" : {
+      "Ios18.expandDims" : 72,
+      "Ios18.mul" : 144,
+      "Ios18.softmax" : 18,
+      "Ios18.matmul" : 36,
+      "Identity" : 1,
+      "Ios16.reduceMean" : 73,
+      "Ios18.greaterEqual" : 1,
+      "Select" : 1,
+      "Ios18.readState" : 37,
+      "Tile" : 36,
+      "Ios18.gather" : 2,
+      "Ios18.add" : 92,
+      "Ios18.layerNorm" : 73,
+      "Ios18.sliceUpdate" : 36,
+      "Ios18.writeState" : 36,
+      "Ios18.reshape" : 110,
+      "Ios18.constexprLutToDense" : 126,
+      "Ios18.conv" : 126,
+      "Ios18.concat" : 108,
+      "Ios18.transpose" : 108,
+      "Ios18.sub" : 73,
+      "Ios18.silu" : 18,
+      "Ios18.sliceByIndex" : 108,
+      "Ios18.squeeze" : 54
+    },
+    "shortDescription" : "Anemll Model: Multifunction FFN+Prefill",
+    "generatedClassName" : "qwen_FFN_PF_lut4_chunk_02of02",
+    "author" : "Converted with Anemll v0.3.0",
+    "modelType" : {
+      "name" : "MLModelType_mlProgram"
+    }
+  }
+]

qwen_FFN_PF_lut4_chunk_02of02.mlmodelc/weights/weight.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f0543f72213b5295cf375c10b60ecb4e04ce02ff7ea2d2edf24def9f47de0ca3
+size 944314880

qwen_embeddings.mlmodelc/analytics/coremldata.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a86e14522df57e2d90d8b058c42c1b8a5f1f6b11186dfb9ed74c390527483570
+size 243

qwen_embeddings.mlmodelc/coremldata.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d04942306ac87412c0358cec9ec357be5b9f781e51e12524e1c763053480e082
+size 501

qwen_embeddings.mlmodelc/metadata.json

@@ -0,0 +1,67 @@
+[
+  {
+    "shortDescription" : "Anemll Model (Embeddings) converted to CoreML",
+    "metadataOutputVersion" : "3.0",
+    "outputSchema" : [
+      {
+        "hasShapeFlexibility" : "0",
+        "isOptional" : "0",
+        "dataType" : "Float16",
+        "formattedType" : "MultiArray (Float16)",
+        "shortDescription" : "",
+        "shape" : "[]",
+        "name" : "hidden_states",
+        "type" : "MultiArray"
+      }
+    ],
+    "version" : "0.3.0",
+    "modelParameters" : [
+
+    ],
+    "author" : "Converted with Anemll v0.3.0",
+    "specificationVersion" : 9,
+    "storagePrecision" : "Float16",
+    "mlProgramOperationTypeHistogram" : {
+      "Ios18.gather" : 1
+    },
+    "computePrecision" : "Mixed (Float16, Int32)",
+    "stateSchema" : [
+
+    ],
+    "isUpdatable" : "0",
+    "availability" : {
+      "macOS" : "15.0",
+      "tvOS" : "18.0",
+      "visionOS" : "2.0",
+      "watchOS" : "11.0",
+      "iOS" : "18.0",
+      "macCatalyst" : "18.0"
+    },
+    "modelType" : {
+      "name" : "MLModelType_mlProgram"
+    },
+    "inputSchema" : [
+      {
+        "shortDescription" : "",
+        "dataType" : "Int32",
+        "hasShapeFlexibility" : "1",
+        "isOptional" : "0",
+        "shapeFlexibility" : "1 × 1 | 1 × 64",
+        "formattedType" : "MultiArray (Int32 1 × 1)",
+        "type" : "MultiArray",
+        "shape" : "[1, 1]",
+        "name" : "input_ids",
+        "enumeratedShapes" : "[[1, 1], [1, 64]]"
+      }
+    ],
+    "userDefinedMetadata" : {
+      "com.anemll.context_length" : "1024",
+      "com.anemll.info" : "Converted with Anemll v0.3.0",
+      "com.github.apple.coremltools.source" : "torch==2.5.0",
+      "com.github.apple.coremltools.version" : "8.3.0",
+      "com.github.apple.coremltools.source_dialect" : "TorchScript"
+    },
+    "generatedClassName" : "qwen_embeddings",
+    "method" : "predict"
+  }
+]

qwen_embeddings.mlmodelc/model.mil

@@ -0,0 +1,11 @@
+program(1.3)
+[buildInfo = dict<string, string>({{"coremlc-component-MIL", "3500.11.1"}, {"coremlc-version", "3500.21.1"}, {"coremltools-component-torch", "2.5.0"}, {"coremltools-source-dialect", "TorchScript"}, {"coremltools-version", "8.3.0"}})]
+{
+    func main<ios18>(tensor<int32, [1, ?]> input_ids) [FlexibleShapeInformation = tuple<tuple<string, dict<string, tensor<int32, [?]>>>, tuple<string, dict<string, dict<string, tensor<int32, [?]>>>>>((("DefaultShapes", {{"input_ids", [1, 1]}}), ("EnumeratedShapes", {{"79ae981e", {{"input_ids", [1, 1]}}}, {"ed9b58c8", {{"input_ids", [1, 64]}}}})))] {
+            int32 hidden_states_axis_0 = const()[name = string("hidden_states_axis_0"), val = int32(0)];
+            int32 hidden_states_batch_dims_0 = const()[name = string("hidden_states_batch_dims_0"), val = int32(0)];
+            bool hidden_states_validate_indices_0 = const()[name = string("hidden_states_validate_indices_0"), val = bool(false)];
+            tensor<fp16, [151936, 2560]> embed_tokens_weight_to_fp16 = const()[name = string("embed_tokens_weight_to_fp16"), val = tensor<fp16, [151936, 2560]>(BLOBFILE(path = string("@model_path/weights/weight.bin"), offset = uint64(64)))];
+            tensor<fp16, [1, ?, 2560]> hidden_states = gather(axis = hidden_states_axis_0, batch_dims = hidden_states_batch_dims_0, indices = input_ids, validate_indices = hidden_states_validate_indices_0, x = embed_tokens_weight_to_fp16)[name = string("hidden_states_cast_fp16")];
+        } -> (hidden_states);
+}

qwen_embeddings.mlmodelc/weights/weight.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:42bfb54b57d5ca84fcaa5a9d0c243febdf0668863d33996cbd93631caf1abfe0
+size 777912448

qwen_lm_head_lut8.mlmodelc/analytics/coremldata.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9ba4088211e9882eeb9250e64cd8dc34128ab2498383a316398e9c988e58f0a2
+size 243

qwen_lm_head_lut8.mlmodelc/coremldata.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5b36cd6b96e663ed87d7c65cfe75f2aab53e1ba16e9cab13f3d0081747bbc80b
+size 898

qwen_lm_head_lut8.mlmodelc/metadata.json

@@ -0,0 +1,220 @@
+[
+  {
+    "shortDescription" : "Anemll Model (LM Head) converted to CoreML",
+    "metadataOutputVersion" : "3.0",
+    "outputSchema" : [
+      {
+        "hasShapeFlexibility" : "0",
+        "isOptional" : "0",
+        "dataType" : "Float16",
+        "formattedType" : "MultiArray (Float16 1 × 1 × 9496)",
+        "shortDescription" : "",
+        "shape" : "[1, 1, 9496]",
+        "name" : "logits1",
+        "type" : "MultiArray"
+      },
+      {
+        "hasShapeFlexibility" : "0",
+        "isOptional" : "0",
+        "dataType" : "Float16",
+        "formattedType" : "MultiArray (Float16 1 × 1 × 9496)",
+        "shortDescription" : "",
+        "shape" : "[1, 1, 9496]",
+        "name" : "logits2",
+        "type" : "MultiArray"
+      },
+      {
+        "hasShapeFlexibility" : "0",
+        "isOptional" : "0",
+        "dataType" : "Float16",
+        "formattedType" : "MultiArray (Float16 1 × 1 × 9496)",
+        "shortDescription" : "",
+        "shape" : "[1, 1, 9496]",
+        "name" : "logits3",
+        "type" : "MultiArray"
+      },
+      {
+        "hasShapeFlexibility" : "0",
+        "isOptional" : "0",
+        "dataType" : "Float16",
+        "formattedType" : "MultiArray (Float16 1 × 1 × 9496)",
+        "shortDescription" : "",
+        "shape" : "[1, 1, 9496]",
+        "name" : "logits4",
+        "type" : "MultiArray"
+      },
+      {
+        "hasShapeFlexibility" : "0",
+        "isOptional" : "0",
+        "dataType" : "Float16",
+        "formattedType" : "MultiArray (Float16 1 × 1 × 9496)",
+        "shortDescription" : "",
+        "shape" : "[1, 1, 9496]",
+        "name" : "logits5",
+        "type" : "MultiArray"
+      },
+      {
+        "hasShapeFlexibility" : "0",
+        "isOptional" : "0",
+        "dataType" : "Float16",
+        "formattedType" : "MultiArray (Float16 1 × 1 × 9496)",
+        "shortDescription" : "",
+        "shape" : "[1, 1, 9496]",
+        "name" : "logits6",
+        "type" : "MultiArray"
+      },
+      {
+        "hasShapeFlexibility" : "0",
+        "isOptional" : "0",
+        "dataType" : "Float16",
+        "formattedType" : "MultiArray (Float16 1 × 1 × 9496)",
+        "shortDescription" : "",
+        "shape" : "[1, 1, 9496]",
+        "name" : "logits7",
+        "type" : "MultiArray"
+      },
+      {
+        "hasShapeFlexibility" : "0",
+        "isOptional" : "0",
+        "dataType" : "Float16",
+        "formattedType" : "MultiArray (Float16 1 × 1 × 9496)",
+        "shortDescription" : "",
+        "shape" : "[1, 1, 9496]",
+        "name" : "logits8",
+        "type" : "MultiArray"
+      },
+      {
+        "hasShapeFlexibility" : "0",
+        "isOptional" : "0",
+        "dataType" : "Float16",
+        "formattedType" : "MultiArray (Float16 1 × 1 × 9496)",
+        "shortDescription" : "",
+        "shape" : "[1, 1, 9496]",
+        "name" : "logits9",
+        "type" : "MultiArray"
+      },
+      {
+        "hasShapeFlexibility" : "0",
+        "isOptional" : "0",
+        "dataType" : "Float16",
+        "formattedType" : "MultiArray (Float16 1 × 1 × 9496)",
+        "shortDescription" : "",
+        "shape" : "[1, 1, 9496]",
+        "name" : "logits10",
+        "type" : "MultiArray"
+      },
+      {
+        "hasShapeFlexibility" : "0",
+        "isOptional" : "0",
+        "dataType" : "Float16",
+        "formattedType" : "MultiArray (Float16 1 × 1 × 9496)",
+        "shortDescription" : "",
+        "shape" : "[1, 1, 9496]",
+        "name" : "logits11",
+        "type" : "MultiArray"
+      },
+      {
+        "hasShapeFlexibility" : "0",
+        "isOptional" : "0",
+        "dataType" : "Float16",
+        "formattedType" : "MultiArray (Float16 1 × 1 × 9496)",
+        "shortDescription" : "",
+        "shape" : "[1, 1, 9496]",
+        "name" : "logits12",
+        "type" : "MultiArray"
+      },
+      {
+        "hasShapeFlexibility" : "0",
+        "isOptional" : "0",
+        "dataType" : "Float16",
+        "formattedType" : "MultiArray (Float16 1 × 1 × 9496)",
+        "shortDescription" : "",
+        "shape" : "[1, 1, 9496]",
+        "name" : "logits13",
+        "type" : "MultiArray"
+      },
+      {
+        "hasShapeFlexibility" : "0",
+        "isOptional" : "0",
+        "dataType" : "Float16",
+        "formattedType" : "MultiArray (Float16 1 × 1 × 9496)",
+        "shortDescription" : "",
+        "shape" : "[1, 1, 9496]",
+        "name" : "logits14",
+        "type" : "MultiArray"
+      },
+      {
+        "hasShapeFlexibility" : "0",
+        "isOptional" : "0",
+        "dataType" : "Float16",
+        "formattedType" : "MultiArray (Float16 1 × 1 × 9496)",
+        "shortDescription" : "",
+        "shape" : "[1, 1, 9496]",
+        "name" : "logits15",
+        "type" : "MultiArray"
+      },
+      {
+        "hasShapeFlexibility" : "0",
+        "isOptional" : "0",
+        "dataType" : "Float16",
+        "formattedType" : "MultiArray (Float16 1 × 1 × 9496)",
+        "shortDescription" : "",
+        "shape" : "[1, 1, 9496]",
+        "name" : "logits16",
+        "type" : "MultiArray"
+      }
+    ],
+    "version" : "0.3.0",
+    "modelParameters" : [
+
+    ],
+    "author" : "Converted with Anemll v0.3.0",
+    "specificationVersion" : 9,
+    "storagePrecision" : "Mixed (Float16, Palettized (19 bits), UInt8)",
+    "mlProgramOperationTypeHistogram" : {
+      "Ios18.transpose" : 17,
+      "Ios18.constexprLutToDense" : 16,
+      "Ios18.expandDims" : 1,
+      "Ios18.conv" : 16,
+      "Ios18.squeeze" : 16
+    },
+    "computePrecision" : "Mixed (Float16, Int32)",
+    "stateSchema" : [
+
+    ],
+    "isUpdatable" : "0",
+    "availability" : {
+      "macOS" : "15.0",
+      "tvOS" : "18.0",
+      "visionOS" : "2.0",
+      "watchOS" : "11.0",
+      "iOS" : "18.0",
+      "macCatalyst" : "18.0"
+    },
+    "modelType" : {
+      "name" : "MLModelType_mlProgram"
+    },
+    "inputSchema" : [
+      {
+        "hasShapeFlexibility" : "0",
+        "isOptional" : "0",
+        "dataType" : "Float16",
+        "formattedType" : "MultiArray (Float16 1 × 1 × 2560)",
+        "shortDescription" : "",
+        "shape" : "[1, 1, 2560]",
+        "name" : "hidden_states",
+        "type" : "MultiArray"
+      }
+    ],
+    "userDefinedMetadata" : {
+      "com.anemll.context_length" : "1024",
+      "com.anemll.info" : "Converted with Anemll v0.3.0",
+      "com.anemll.lut_bits" : "8",
+      "com.github.apple.coremltools.source" : "torch==2.5.0",
+      "com.github.apple.coremltools.version" : "8.3.0",
+      "com.github.apple.coremltools.source_dialect" : "TorchScript"
+    },
+    "generatedClassName" : "qwen_lm_head_lut8",
+    "method" : "predict"
+  }
+]

qwen_lm_head_lut8.mlmodelc/model.mil

@@ -0,0 +1,186 @@
+program(1.3)
+[buildInfo = dict<string, string>({{"coremlc-component-MIL", "3500.11.1"}, {"coremlc-version", "3500.21.1"}})]
+{
+    func main<ios18>(tensor<fp16, [1, 1, 2560]> hidden_states) {
+            tensor<int32, [3]> var_5 = const()[name = string("op_5"), val = tensor<int32, [3]>([0, 2, 1])];
+            tensor<int32, [1]> input_axes_0 = const()[name = string("input_axes_0"), val = tensor<int32, [1]>([2])];
+            tensor<fp16, [1, 2560, 1]> var_6_cast_fp16 = transpose(perm = var_5, x = hidden_states)[name = string("transpose_16")];
+            tensor<fp16, [1, 2560, 1, 1]> input_cast_fp16 = expand_dims(axes = input_axes_0, x = var_6_cast_fp16)[name = string("input_cast_fp16")];
+            string var_29_pad_type_0 = const()[name = string("op_29_pad_type_0"), val = string("valid")];
+            tensor<int32, [2]> var_29_strides_0 = const()[name = string("op_29_strides_0"), val = tensor<int32, [2]>([1, 1])];
+            tensor<int32, [4]> var_29_pad_0 = const()[name = string("op_29_pad_0"), val = tensor<int32, [4]>([0, 0, 0, 0])];
+            tensor<int32, [2]> var_29_dilations_0 = const()[name = string("op_29_dilations_0"), val = tensor<int32, [2]>([1, 1])];
+            int32 var_29_groups_0 = const()[name = string("op_29_groups_0"), val = int32(1)];
+            tensor<fp16, [9496, 2560, 1, 1]> op_9_promoted_to_fp16_palettized = constexpr_lut_to_dense(indices = tensor<uint8, [9496, 2560, 1, 1]>(BLOBFILE(path = string("@model_path/weights/weight.bin"), offset = uint64(64))), lut = tensor<fp16, [1187, 1, 1, 1, 256, 1]>(BLOBFILE(path = string("@model_path/weights/weight.bin"), offset = uint64(24309888))))[name = string("op_9_promoted_to_fp16_palettized")];
+            tensor<fp16, [1, 9496, 1, 1]> var_29_cast_fp16 = conv(dilations = var_29_dilations_0, groups = var_29_groups_0, pad = var_29_pad_0, pad_type = var_29_pad_type_0, strides = var_29_strides_0, weight = op_9_promoted_to_fp16_palettized, x = input_cast_fp16)[name = string("op_29_cast_fp16")];
+            tensor<int32, [1]> var_31_axes_0 = const()[name = string("op_31_axes_0"), val = tensor<int32, [1]>([2])];
+            tensor<fp16, [1, 9496, 1]> var_31_cast_fp16 = squeeze(axes = var_31_axes_0, x = var_29_cast_fp16)[name = string("op_31_cast_fp16")];
+            tensor<int32, [3]> var_34_perm_0 = const()[name = string("op_34_perm_0"), val = tensor<int32, [3]>([0, 2, 1])];
+            string var_55_pad_type_0 = const()[name = string("op_55_pad_type_0"), val = string("valid")];
+            tensor<int32, [2]> var_55_strides_0 = const()[name = string("op_55_strides_0"), val = tensor<int32, [2]>([1, 1])];
+            tensor<int32, [4]> var_55_pad_0 = const()[name = string("op_55_pad_0"), val = tensor<int32, [4]>([0, 0, 0, 0])];
+            tensor<int32, [2]> var_55_dilations_0 = const()[name = string("op_55_dilations_0"), val = tensor<int32, [2]>([1, 1])];
+            int32 var_55_groups_0 = const()[name = string("op_55_groups_0"), val = int32(1)];
+            tensor<fp16, [9496, 2560, 1, 1]> op_35_promoted_to_fp16_palettized = constexpr_lut_to_dense(indices = tensor<uint8, [9496, 2560, 1, 1]>(BLOBFILE(path = string("@model_path/weights/weight.bin"), offset = uint64(24917696))), lut = tensor<fp16, [1187, 1, 1, 1, 256, 1]>(BLOBFILE(path = string("@model_path/weights/weight.bin"), offset = uint64(49227520))))[name = string("op_35_promoted_to_fp16_palettized")];
+            tensor<fp16, [1, 9496, 1, 1]> var_55_cast_fp16 = conv(dilations = var_55_dilations_0, groups = var_55_groups_0, pad = var_55_pad_0, pad_type = var_55_pad_type_0, strides = var_55_strides_0, weight = op_35_promoted_to_fp16_palettized, x = input_cast_fp16)[name = string("op_55_cast_fp16")];
+            tensor<int32, [1]> var_57_axes_0 = const()[name = string("op_57_axes_0"), val = tensor<int32, [1]>([2])];
+            tensor<fp16, [1, 9496, 1]> var_57_cast_fp16 = squeeze(axes = var_57_axes_0, x = var_55_cast_fp16)[name = string("op_57_cast_fp16")];
+            tensor<int32, [3]> var_60_perm_0 = const()[name = string("op_60_perm_0"), val = tensor<int32, [3]>([0, 2, 1])];
+            string var_81_pad_type_0 = const()[name = string("op_81_pad_type_0"), val = string("valid")];
+            tensor<int32, [2]> var_81_strides_0 = const()[name = string("op_81_strides_0"), val = tensor<int32, [2]>([1, 1])];
+            tensor<int32, [4]> var_81_pad_0 = const()[name = string("op_81_pad_0"), val = tensor<int32, [4]>([0, 0, 0, 0])];
+            tensor<int32, [2]> var_81_dilations_0 = const()[name = string("op_81_dilations_0"), val = tensor<int32, [2]>([1, 1])];
+            int32 var_81_groups_0 = const()[name = string("op_81_groups_0"), val = int32(1)];
+            tensor<fp16, [9496, 2560, 1, 1]> op_61_promoted_to_fp16_palettized = constexpr_lut_to_dense(indices = tensor<uint8, [9496, 2560, 1, 1]>(BLOBFILE(path = string("@model_path/weights/weight.bin"), offset = uint64(49835328))), lut = tensor<fp16, [1187, 1, 1, 1, 256, 1]>(BLOBFILE(path = string("@model_path/weights/weight.bin"), offset = uint64(74145152))))[name = string("op_61_promoted_to_fp16_palettized")];
+            tensor<fp16, [1, 9496, 1, 1]> var_81_cast_fp16 = conv(dilations = var_81_dilations_0, groups = var_81_groups_0, pad = var_81_pad_0, pad_type = var_81_pad_type_0, strides = var_81_strides_0, weight = op_61_promoted_to_fp16_palettized, x = input_cast_fp16)[name = string("op_81_cast_fp16")];
+            tensor<int32, [1]> var_83_axes_0 = const()[name = string("op_83_axes_0"), val = tensor<int32, [1]>([2])];
+            tensor<fp16, [1, 9496, 1]> var_83_cast_fp16 = squeeze(axes = var_83_axes_0, x = var_81_cast_fp16)[name = string("op_83_cast_fp16")];
+            tensor<int32, [3]> var_86_perm_0 = const()[name = string("op_86_perm_0"), val = tensor<int32, [3]>([0, 2, 1])];
+            string var_107_pad_type_0 = const()[name = string("op_107_pad_type_0"), val = string("valid")];
+            tensor<int32, [2]> var_107_strides_0 = const()[name = string("op_107_strides_0"), val = tensor<int32, [2]>([1, 1])];
+            tensor<int32, [4]> var_107_pad_0 = const()[name = string("op_107_pad_0"), val = tensor<int32, [4]>([0, 0, 0, 0])];
+            tensor<int32, [2]> var_107_dilations_0 = const()[name = string("op_107_dilations_0"), val = tensor<int32, [2]>([1, 1])];
+            int32 var_107_groups_0 = const()[name = string("op_107_groups_0"), val = int32(1)];
+            tensor<fp16, [9496, 2560, 1, 1]> op_87_promoted_to_fp16_palettized = constexpr_lut_to_dense(indices = tensor<uint8, [9496, 2560, 1, 1]>(BLOBFILE(path = string("@model_path/weights/weight.bin"), offset = uint64(74752960))), lut = tensor<fp16, [1187, 1, 1, 1, 256, 1]>(BLOBFILE(path = string("@model_path/weights/weight.bin"), offset = uint64(99062784))))[name = string("op_87_promoted_to_fp16_palettized")];
+            tensor<fp16, [1, 9496, 1, 1]> var_107_cast_fp16 = conv(dilations = var_107_dilations_0, groups = var_107_groups_0, pad = var_107_pad_0, pad_type = var_107_pad_type_0, strides = var_107_strides_0, weight = op_87_promoted_to_fp16_palettized, x = input_cast_fp16)[name = string("op_107_cast_fp16")];
+            tensor<int32, [1]> var_109_axes_0 = const()[name = string("op_109_axes_0"), val = tensor<int32, [1]>([2])];
+            tensor<fp16, [1, 9496, 1]> var_109_cast_fp16 = squeeze(axes = var_109_axes_0, x = var_107_cast_fp16)[name = string("op_109_cast_fp16")];
+            tensor<int32, [3]> var_112_perm_0 = const()[name = string("op_112_perm_0"), val = tensor<int32, [3]>([0, 2, 1])];
+            string var_133_pad_type_0 = const()[name = string("op_133_pad_type_0"), val = string("valid")];
+            tensor<int32, [2]> var_133_strides_0 = const()[name = string("op_133_strides_0"), val = tensor<int32, [2]>([1, 1])];
+            tensor<int32, [4]> var_133_pad_0 = const()[name = string("op_133_pad_0"), val = tensor<int32, [4]>([0, 0, 0, 0])];
+            tensor<int32, [2]> var_133_dilations_0 = const()[name = string("op_133_dilations_0"), val = tensor<int32, [2]>([1, 1])];
+            int32 var_133_groups_0 = const()[name = string("op_133_groups_0"), val = int32(1)];
+            tensor<fp16, [9496, 2560, 1, 1]> op_113_promoted_to_fp16_palettized = constexpr_lut_to_dense(indices = tensor<uint8, [9496, 2560, 1, 1]>(BLOBFILE(path = string("@model_path/weights/weight.bin"), offset = uint64(99670592))), lut = tensor<fp16, [1187, 1, 1, 1, 256, 1]>(BLOBFILE(path = string("@model_path/weights/weight.bin"), offset = uint64(123980416))))[name = string("op_113_promoted_to_fp16_palettized")];
+            tensor<fp16, [1, 9496, 1, 1]> var_133_cast_fp16 = conv(dilations = var_133_dilations_0, groups = var_133_groups_0, pad = var_133_pad_0, pad_type = var_133_pad_type_0, strides = var_133_strides_0, weight = op_113_promoted_to_fp16_palettized, x = input_cast_fp16)[name = string("op_133_cast_fp16")];
+            tensor<int32, [1]> var_135_axes_0 = const()[name = string("op_135_axes_0"), val = tensor<int32, [1]>([2])];
+            tensor<fp16, [1, 9496, 1]> var_135_cast_fp16 = squeeze(axes = var_135_axes_0, x = var_133_cast_fp16)[name = string("op_135_cast_fp16")];
+            tensor<int32, [3]> var_138_perm_0 = const()[name = string("op_138_perm_0"), val = tensor<int32, [3]>([0, 2, 1])];
+            string var_159_pad_type_0 = const()[name = string("op_159_pad_type_0"), val = string("valid")];
+            tensor<int32, [2]> var_159_strides_0 = const()[name = string("op_159_strides_0"), val = tensor<int32, [2]>([1, 1])];
+            tensor<int32, [4]> var_159_pad_0 = const()[name = string("op_159_pad_0"), val = tensor<int32, [4]>([0, 0, 0, 0])];
+            tensor<int32, [2]> var_159_dilations_0 = const()[name = string("op_159_dilations_0"), val = tensor<int32, [2]>([1, 1])];
+            int32 var_159_groups_0 = const()[name = string("op_159_groups_0"), val = int32(1)];
+            tensor<fp16, [9496, 2560, 1, 1]> op_139_promoted_to_fp16_palettized = constexpr_lut_to_dense(indices = tensor<uint8, [9496, 2560, 1, 1]>(BLOBFILE(path = string("@model_path/weights/weight.bin"), offset = uint64(124588224))), lut = tensor<fp16, [1187, 1, 1, 1, 256, 1]>(BLOBFILE(path = string("@model_path/weights/weight.bin"), offset = uint64(148898048))))[name = string("op_139_promoted_to_fp16_palettized")];
+            tensor<fp16, [1, 9496, 1, 1]> var_159_cast_fp16 = conv(dilations = var_159_dilations_0, groups = var_159_groups_0, pad = var_159_pad_0, pad_type = var_159_pad_type_0, strides = var_159_strides_0, weight = op_139_promoted_to_fp16_palettized, x = input_cast_fp16)[name = string("op_159_cast_fp16")];
+            tensor<int32, [1]> var_161_axes_0 = const()[name = string("op_161_axes_0"), val = tensor<int32, [1]>([2])];
+            tensor<fp16, [1, 9496, 1]> var_161_cast_fp16 = squeeze(axes = var_161_axes_0, x = var_159_cast_fp16)[name = string("op_161_cast_fp16")];
+            tensor<int32, [3]> var_164_perm_0 = const()[name = string("op_164_perm_0"), val = tensor<int32, [3]>([0, 2, 1])];
+            string var_185_pad_type_0 = const()[name = string("op_185_pad_type_0"), val = string("valid")];
+            tensor<int32, [2]> var_185_strides_0 = const()[name = string("op_185_strides_0"), val = tensor<int32, [2]>([1, 1])];
+            tensor<int32, [4]> var_185_pad_0 = const()[name = string("op_185_pad_0"), val = tensor<int32, [4]>([0, 0, 0, 0])];
+            tensor<int32, [2]> var_185_dilations_0 = const()[name = string("op_185_dilations_0"), val = tensor<int32, [2]>([1, 1])];
+            int32 var_185_groups_0 = const()[name = string("op_185_groups_0"), val = int32(1)];
+            tensor<fp16, [9496, 2560, 1, 1]> op_165_promoted_to_fp16_palettized = constexpr_lut_to_dense(indices = tensor<uint8, [9496, 2560, 1, 1]>(BLOBFILE(path = string("@model_path/weights/weight.bin"), offset = uint64(149505856))), lut = tensor<fp16, [1187, 1, 1, 1, 256, 1]>(BLOBFILE(path = string("@model_path/weights/weight.bin"), offset = uint64(173815680))))[name = string("op_165_promoted_to_fp16_palettized")];
+            tensor<fp16, [1, 9496, 1, 1]> var_185_cast_fp16 = conv(dilations = var_185_dilations_0, groups = var_185_groups_0, pad = var_185_pad_0, pad_type = var_185_pad_type_0, strides = var_185_strides_0, weight = op_165_promoted_to_fp16_palettized, x = input_cast_fp16)[name = string("op_185_cast_fp16")];
+            tensor<int32, [1]> var_187_axes_0 = const()[name = string("op_187_axes_0"), val = tensor<int32, [1]>([2])];
+            tensor<fp16, [1, 9496, 1]> var_187_cast_fp16 = squeeze(axes = var_187_axes_0, x = var_185_cast_fp16)[name = string("op_187_cast_fp16")];
+            tensor<int32, [3]> var_190_perm_0 = const()[name = string("op_190_perm_0"), val = tensor<int32, [3]>([0, 2, 1])];
+            string var_211_pad_type_0 = const()[name = string("op_211_pad_type_0"), val = string("valid")];
+            tensor<int32, [2]> var_211_strides_0 = const()[name = string("op_211_strides_0"), val = tensor<int32, [2]>([1, 1])];
+            tensor<int32, [4]> var_211_pad_0 = const()[name = string("op_211_pad_0"), val = tensor<int32, [4]>([0, 0, 0, 0])];
+            tensor<int32, [2]> var_211_dilations_0 = const()[name = string("op_211_dilations_0"), val = tensor<int32, [2]>([1, 1])];
+            int32 var_211_groups_0 = const()[name = string("op_211_groups_0"), val = int32(1)];
+            tensor<fp16, [9496, 2560, 1, 1]> op_191_promoted_to_fp16_palettized = constexpr_lut_to_dense(indices = tensor<uint8, [9496, 2560, 1, 1]>(BLOBFILE(path = string("@model_path/weights/weight.bin"), offset = uint64(174423488))), lut = tensor<fp16, [1187, 1, 1, 1, 256, 1]>(BLOBFILE(path = string("@model_path/weights/weight.bin"), offset = uint64(198733312))))[name = string("op_191_promoted_to_fp16_palettized")];
+            tensor<fp16, [1, 9496, 1, 1]> var_211_cast_fp16 = conv(dilations = var_211_dilations_0, groups = var_211_groups_0, pad = var_211_pad_0, pad_type = var_211_pad_type_0, strides = var_211_strides_0, weight = op_191_promoted_to_fp16_palettized, x = input_cast_fp16)[name = string("op_211_cast_fp16")];
+            tensor<int32, [1]> var_213_axes_0 = const()[name = string("op_213_axes_0"), val = tensor<int32, [1]>([2])];
+            tensor<fp16, [1, 9496, 1]> var_213_cast_fp16 = squeeze(axes = var_213_axes_0, x = var_211_cast_fp16)[name = string("op_213_cast_fp16")];
+            tensor<int32, [3]> var_216_perm_0 = const()[name = string("op_216_perm_0"), val = tensor<int32, [3]>([0, 2, 1])];
+            string var_237_pad_type_0 = const()[name = string("op_237_pad_type_0"), val = string("valid")];
+            tensor<int32, [2]> var_237_strides_0 = const()[name = string("op_237_strides_0"), val = tensor<int32, [2]>([1, 1])];
+            tensor<int32, [4]> var_237_pad_0 = const()[name = string("op_237_pad_0"), val = tensor<int32, [4]>([0, 0, 0, 0])];
+            tensor<int32, [2]> var_237_dilations_0 = const()[name = string("op_237_dilations_0"), val = tensor<int32, [2]>([1, 1])];
+            int32 var_237_groups_0 = const()[name = string("op_237_groups_0"), val = int32(1)];
+            tensor<fp16, [9496, 2560, 1, 1]> op_217_promoted_to_fp16_palettized = constexpr_lut_to_dense(indices = tensor<uint8, [9496, 2560, 1, 1]>(BLOBFILE(path = string("@model_path/weights/weight.bin"), offset = uint64(199341120))), lut = tensor<fp16, [1187, 1, 1, 1, 256, 1]>(BLOBFILE(path = string("@model_path/weights/weight.bin"), offset = uint64(223650944))))[name = string("op_217_promoted_to_fp16_palettized")];
+            tensor<fp16, [1, 9496, 1, 1]> var_237_cast_fp16 = conv(dilations = var_237_dilations_0, groups = var_237_groups_0, pad = var_237_pad_0, pad_type = var_237_pad_type_0, strides = var_237_strides_0, weight = op_217_promoted_to_fp16_palettized, x = input_cast_fp16)[name = string("op_237_cast_fp16")];
+            tensor<int32, [1]> var_239_axes_0 = const()[name = string("op_239_axes_0"), val = tensor<int32, [1]>([2])];
+            tensor<fp16, [1, 9496, 1]> var_239_cast_fp16 = squeeze(axes = var_239_axes_0, x = var_237_cast_fp16)[name = string("op_239_cast_fp16")];
+            tensor<int32, [3]> var_242_perm_0 = const()[name = string("op_242_perm_0"), val = tensor<int32, [3]>([0, 2, 1])];
+            string var_263_pad_type_0 = const()[name = string("op_263_pad_type_0"), val = string("valid")];
+            tensor<int32, [2]> var_263_strides_0 = const()[name = string("op_263_strides_0"), val = tensor<int32, [2]>([1, 1])];
+            tensor<int32, [4]> var_263_pad_0 = const()[name = string("op_263_pad_0"), val = tensor<int32, [4]>([0, 0, 0, 0])];
+            tensor<int32, [2]> var_263_dilations_0 = const()[name = string("op_263_dilations_0"), val = tensor<int32, [2]>([1, 1])];
+            int32 var_263_groups_0 = const()[name = string("op_263_groups_0"), val = int32(1)];
+            tensor<fp16, [9496, 2560, 1, 1]> op_243_promoted_to_fp16_palettized = constexpr_lut_to_dense(indices = tensor<uint8, [9496, 2560, 1, 1]>(BLOBFILE(path = string("@model_path/weights/weight.bin"), offset = uint64(224258752))), lut = tensor<fp16, [1187, 1, 1, 1, 256, 1]>(BLOBFILE(path = string("@model_path/weights/weight.bin"), offset = uint64(248568576))))[name = string("op_243_promoted_to_fp16_palettized")];
+            tensor<fp16, [1, 9496, 1, 1]> var_263_cast_fp16 = conv(dilations = var_263_dilations_0, groups = var_263_groups_0, pad = var_263_pad_0, pad_type = var_263_pad_type_0, strides = var_263_strides_0, weight = op_243_promoted_to_fp16_palettized, x = input_cast_fp16)[name = string("op_263_cast_fp16")];
+            tensor<int32, [1]> var_265_axes_0 = const()[name = string("op_265_axes_0"), val = tensor<int32, [1]>([2])];
+            tensor<fp16, [1, 9496, 1]> var_265_cast_fp16 = squeeze(axes = var_265_axes_0, x = var_263_cast_fp16)[name = string("op_265_cast_fp16")];
+            tensor<int32, [3]> var_268_perm_0 = const()[name = string("op_268_perm_0"), val = tensor<int32, [3]>([0, 2, 1])];
+            string var_289_pad_type_0 = const()[name = string("op_289_pad_type_0"), val = string("valid")];
+            tensor<int32, [2]> var_289_strides_0 = const()[name = string("op_289_strides_0"), val = tensor<int32, [2]>([1, 1])];
+            tensor<int32, [4]> var_289_pad_0 = const()[name = string("op_289_pad_0"), val = tensor<int32, [4]>([0, 0, 0, 0])];
+            tensor<int32, [2]> var_289_dilations_0 = const()[name = string("op_289_dilations_0"), val = tensor<int32, [2]>([1, 1])];
+            int32 var_289_groups_0 = const()[name = string("op_289_groups_0"), val = int32(1)];
+            tensor<fp16, [9496, 2560, 1, 1]> op_269_promoted_to_fp16_palettized = constexpr_lut_to_dense(indices = tensor<uint8, [9496, 2560, 1, 1]>(BLOBFILE(path = string("@model_path/weights/weight.bin"), offset = uint64(249176384))), lut = tensor<fp16, [1187, 1, 1, 1, 256, 1]>(BLOBFILE(path = string("@model_path/weights/weight.bin"), offset = uint64(273486208))))[name = string("op_269_promoted_to_fp16_palettized")];
+            tensor<fp16, [1, 9496, 1, 1]> var_289_cast_fp16 = conv(dilations = var_289_dilations_0, groups = var_289_groups_0, pad = var_289_pad_0, pad_type = var_289_pad_type_0, strides = var_289_strides_0, weight = op_269_promoted_to_fp16_palettized, x = input_cast_fp16)[name = string("op_289_cast_fp16")];
+            tensor<int32, [1]> var_291_axes_0 = const()[name = string("op_291_axes_0"), val = tensor<int32, [1]>([2])];
+            tensor<fp16, [1, 9496, 1]> var_291_cast_fp16 = squeeze(axes = var_291_axes_0, x = var_289_cast_fp16)[name = string("op_291_cast_fp16")];
+            tensor<int32, [3]> var_294_perm_0 = const()[name = string("op_294_perm_0"), val = tensor<int32, [3]>([0, 2, 1])];
+            string var_315_pad_type_0 = const()[name = string("op_315_pad_type_0"), val = string("valid")];
+            tensor<int32, [2]> var_315_strides_0 = const()[name = string("op_315_strides_0"), val = tensor<int32, [2]>([1, 1])];
+            tensor<int32, [4]> var_315_pad_0 = const()[name = string("op_315_pad_0"), val = tensor<int32, [4]>([0, 0, 0, 0])];
+            tensor<int32, [2]> var_315_dilations_0 = const()[name = string("op_315_dilations_0"), val = tensor<int32, [2]>([1, 1])];
+            int32 var_315_groups_0 = const()[name = string("op_315_groups_0"), val = int32(1)];
+            tensor<fp16, [9496, 2560, 1, 1]> op_295_promoted_to_fp16_palettized = constexpr_lut_to_dense(indices = tensor<uint8, [9496, 2560, 1, 1]>(BLOBFILE(path = string("@model_path/weights/weight.bin"), offset = uint64(274094016))), lut = tensor<fp16, [1187, 1, 1, 1, 256, 1]>(BLOBFILE(path = string("@model_path/weights/weight.bin"), offset = uint64(298403840))))[name = string("op_295_promoted_to_fp16_palettized")];
+            tensor<fp16, [1, 9496, 1, 1]> var_315_cast_fp16 = conv(dilations = var_315_dilations_0, groups = var_315_groups_0, pad = var_315_pad_0, pad_type = var_315_pad_type_0, strides = var_315_strides_0, weight = op_295_promoted_to_fp16_palettized, x = input_cast_fp16)[name = string("op_315_cast_fp16")];
+            tensor<int32, [1]> var_317_axes_0 = const()[name = string("op_317_axes_0"), val = tensor<int32, [1]>([2])];
+            tensor<fp16, [1, 9496, 1]> var_317_cast_fp16 = squeeze(axes = var_317_axes_0, x = var_315_cast_fp16)[name = string("op_317_cast_fp16")];
+            tensor<int32, [3]> var_320_perm_0 = const()[name = string("op_320_perm_0"), val = tensor<int32, [3]>([0, 2, 1])];
+            string var_341_pad_type_0 = const()[name = string("op_341_pad_type_0"), val = string("valid")];
+            tensor<int32, [2]> var_341_strides_0 = const()[name = string("op_341_strides_0"), val = tensor<int32, [2]>([1, 1])];
+            tensor<int32, [4]> var_341_pad_0 = const()[name = string("op_341_pad_0"), val = tensor<int32, [4]>([0, 0, 0, 0])];
+            tensor<int32, [2]> var_341_dilations_0 = const()[name = string("op_341_dilations_0"), val = tensor<int32, [2]>([1, 1])];
+            int32 var_341_groups_0 = const()[name = string("op_341_groups_0"), val = int32(1)];
+            tensor<fp16, [9496, 2560, 1, 1]> op_321_promoted_to_fp16_palettized = constexpr_lut_to_dense(indices = tensor<uint8, [9496, 2560, 1, 1]>(BLOBFILE(path = string("@model_path/weights/weight.bin"), offset = uint64(299011648))), lut = tensor<fp16, [1187, 1, 1, 1, 256, 1]>(BLOBFILE(path = string("@model_path/weights/weight.bin"), offset = uint64(323321472))))[name = string("op_321_promoted_to_fp16_palettized")];
+            tensor<fp16, [1, 9496, 1, 1]> var_341_cast_fp16 = conv(dilations = var_341_dilations_0, groups = var_341_groups_0, pad = var_341_pad_0, pad_type = var_341_pad_type_0, strides = var_341_strides_0, weight = op_321_promoted_to_fp16_palettized, x = input_cast_fp16)[name = string("op_341_cast_fp16")];
+            tensor<int32, [1]> var_343_axes_0 = const()[name = string("op_343_axes_0"), val = tensor<int32, [1]>([2])];
+            tensor<fp16, [1, 9496, 1]> var_343_cast_fp16 = squeeze(axes = var_343_axes_0, x = var_341_cast_fp16)[name = string("op_343_cast_fp16")];
+            tensor<int32, [3]> var_346_perm_0 = const()[name = string("op_346_perm_0"), val = tensor<int32, [3]>([0, 2, 1])];
+            string var_367_pad_type_0 = const()[name = string("op_367_pad_type_0"), val = string("valid")];
+            tensor<int32, [2]> var_367_strides_0 = const()[name = string("op_367_strides_0"), val = tensor<int32, [2]>([1, 1])];
+            tensor<int32, [4]> var_367_pad_0 = const()[name = string("op_367_pad_0"), val = tensor<int32, [4]>([0, 0, 0, 0])];
+            tensor<int32, [2]> var_367_dilations_0 = const()[name = string("op_367_dilations_0"), val = tensor<int32, [2]>([1, 1])];
+            int32 var_367_groups_0 = const()[name = string("op_367_groups_0"), val = int32(1)];
+            tensor<fp16, [9496, 2560, 1, 1]> op_347_promoted_to_fp16_palettized = constexpr_lut_to_dense(indices = tensor<uint8, [9496, 2560, 1, 1]>(BLOBFILE(path = string("@model_path/weights/weight.bin"), offset = uint64(323929280))), lut = tensor<fp16, [1187, 1, 1, 1, 256, 1]>(BLOBFILE(path = string("@model_path/weights/weight.bin"), offset = uint64(348239104))))[name = string("op_347_promoted_to_fp16_palettized")];
+            tensor<fp16, [1, 9496, 1, 1]> var_367_cast_fp16 = conv(dilations = var_367_dilations_0, groups = var_367_groups_0, pad = var_367_pad_0, pad_type = var_367_pad_type_0, strides = var_367_strides_0, weight = op_347_promoted_to_fp16_palettized, x = input_cast_fp16)[name = string("op_367_cast_fp16")];
+            tensor<int32, [1]> var_369_axes_0 = const()[name = string("op_369_axes_0"), val = tensor<int32, [1]>([2])];
+            tensor<fp16, [1, 9496, 1]> var_369_cast_fp16 = squeeze(axes = var_369_axes_0, x = var_367_cast_fp16)[name = string("op_369_cast_fp16")];
+            tensor<int32, [3]> var_372_perm_0 = const()[name = string("op_372_perm_0"), val = tensor<int32, [3]>([0, 2, 1])];
+            string var_393_pad_type_0 = const()[name = string("op_393_pad_type_0"), val = string("valid")];
+            tensor<int32, [2]> var_393_strides_0 = const()[name = string("op_393_strides_0"), val = tensor<int32, [2]>([1, 1])];
+            tensor<int32, [4]> var_393_pad_0 = const()[name = string("op_393_pad_0"), val = tensor<int32, [4]>([0, 0, 0, 0])];
+            tensor<int32, [2]> var_393_dilations_0 = const()[name = string("op_393_dilations_0"), val = tensor<int32, [2]>([1, 1])];
+            int32 var_393_groups_0 = const()[name = string("op_393_groups_0"), val = int32(1)];
+            tensor<fp16, [9496, 2560, 1, 1]> op_373_promoted_to_fp16_palettized = constexpr_lut_to_dense(indices = tensor<uint8, [9496, 2560, 1, 1]>(BLOBFILE(path = string("@model_path/weights/weight.bin"), offset = uint64(348846912))), lut = tensor<fp16, [1187, 1, 1, 1, 256, 1]>(BLOBFILE(path = string("@model_path/weights/weight.bin"), offset = uint64(373156736))))[name = string("op_373_promoted_to_fp16_palettized")];
+            tensor<fp16, [1, 9496, 1, 1]> var_393_cast_fp16 = conv(dilations = var_393_dilations_0, groups = var_393_groups_0, pad = var_393_pad_0, pad_type = var_393_pad_type_0, strides = var_393_strides_0, weight = op_373_promoted_to_fp16_palettized, x = input_cast_fp16)[name = string("op_393_cast_fp16")];
+            tensor<int32, [1]> var_395_axes_0 = const()[name = string("op_395_axes_0"), val = tensor<int32, [1]>([2])];
+            tensor<fp16, [1, 9496, 1]> var_395_cast_fp16 = squeeze(axes = var_395_axes_0, x = var_393_cast_fp16)[name = string("op_395_cast_fp16")];
+            tensor<int32, [3]> var_398_perm_0 = const()[name = string("op_398_perm_0"), val = tensor<int32, [3]>([0, 2, 1])];
+            string var_419_pad_type_0 = const()[name = string("op_419_pad_type_0"), val = string("valid")];
+            tensor<int32, [2]> var_419_strides_0 = const()[name = string("op_419_strides_0"), val = tensor<int32, [2]>([1, 1])];
+            tensor<int32, [4]> var_419_pad_0 = const()[name = string("op_419_pad_0"), val = tensor<int32, [4]>([0, 0, 0, 0])];
+            tensor<int32, [2]> var_419_dilations_0 = const()[name = string("op_419_dilations_0"), val = tensor<int32, [2]>([1, 1])];
+            int32 var_419_groups_0 = const()[name = string("op_419_groups_0"), val = int32(1)];
+            tensor<fp16, [9496, 2560, 1, 1]> op_399_promoted_to_fp16_palettized = constexpr_lut_to_dense(indices = tensor<uint8, [9496, 2560, 1, 1]>(BLOBFILE(path = string("@model_path/weights/weight.bin"), offset = uint64(373764544))), lut = tensor<fp16, [1187, 1, 1, 1, 256, 1]>(BLOBFILE(path = string("@model_path/weights/weight.bin"), offset = uint64(398074368))))[name = string("op_399_promoted_to_fp16_palettized")];
+            tensor<fp16, [1, 9496, 1, 1]> var_419_cast_fp16 = conv(dilations = var_419_dilations_0, groups = var_419_groups_0, pad = var_419_pad_0, pad_type = var_419_pad_type_0, strides = var_419_strides_0, weight = op_399_promoted_to_fp16_palettized, x = input_cast_fp16)[name = string("op_419_cast_fp16")];
+            tensor<int32, [1]> var_421_axes_0 = const()[name = string("op_421_axes_0"), val = tensor<int32, [1]>([2])];
+            tensor<fp16, [1, 9496, 1]> var_421_cast_fp16 = squeeze(axes = var_421_axes_0, x = var_419_cast_fp16)[name = string("op_421_cast_fp16")];
+            tensor<int32, [3]> var_424_perm_0 = const()[name = string("op_424_perm_0"), val = tensor<int32, [3]>([0, 2, 1])];
+            tensor<fp16, [1, 1, 9496]> logits1 = transpose(perm = var_34_perm_0, x = var_31_cast_fp16)[name = string("transpose_0")];
+            tensor<fp16, [1, 1, 9496]> logits2 = transpose(perm = var_60_perm_0, x = var_57_cast_fp16)[name = string("transpose_1")];
+            tensor<fp16, [1, 1, 9496]> logits3 = transpose(perm = var_86_perm_0, x = var_83_cast_fp16)[name = string("transpose_2")];
+            tensor<fp16, [1, 1, 9496]> logits4 = transpose(perm = var_112_perm_0, x = var_109_cast_fp16)[name = string("transpose_3")];
+            tensor<fp16, [1, 1, 9496]> logits5 = transpose(perm = var_138_perm_0, x = var_135_cast_fp16)[name = string("transpose_4")];
+            tensor<fp16, [1, 1, 9496]> logits6 = transpose(perm = var_164_perm_0, x = var_161_cast_fp16)[name = string("transpose_5")];
+            tensor<fp16, [1, 1, 9496]> logits7 = transpose(perm = var_190_perm_0, x = var_187_cast_fp16)[name = string("transpose_6")];
+            tensor<fp16, [1, 1, 9496]> logits8 = transpose(perm = var_216_perm_0, x = var_213_cast_fp16)[name = string("transpose_7")];
+            tensor<fp16, [1, 1, 9496]> logits9 = transpose(perm = var_242_perm_0, x = var_239_cast_fp16)[name = string("transpose_8")];
+            tensor<fp16, [1, 1, 9496]> logits10 = transpose(perm = var_268_perm_0, x = var_265_cast_fp16)[name = string("transpose_9")];
+            tensor<fp16, [1, 1, 9496]> logits11 = transpose(perm = var_294_perm_0, x = var_291_cast_fp16)[name = string("transpose_10")];
+            tensor<fp16, [1, 1, 9496]> logits12 = transpose(perm = var_320_perm_0, x = var_317_cast_fp16)[name = string("transpose_11")];
+            tensor<fp16, [1, 1, 9496]> logits13 = transpose(perm = var_346_perm_0, x = var_343_cast_fp16)[name = string("transpose_12")];
+            tensor<fp16, [1, 1, 9496]> logits14 = transpose(perm = var_372_perm_0, x = var_369_cast_fp16)[name = string("transpose_13")];
+            tensor<fp16, [1, 1, 9496]> logits15 = transpose(perm = var_398_perm_0, x = var_395_cast_fp16)[name = string("transpose_14")];
+            tensor<fp16, [1, 1, 9496]> logits16 = transpose(perm = var_424_perm_0, x = var_421_cast_fp16)[name = string("transpose_15")];
+        } -> (logits1, logits2, logits3, logits4, logits5, logits6, logits7, logits8, logits9, logits10, logits11, logits12, logits13, logits14, logits15, logits16);
+}

qwen_lm_head_lut8.mlmodelc/weights/weight.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:84dced6345d9f025b22105d74af6176fafae44ba69644668d07c364dc81a22c8
+size 398682176

tokenizer.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:aeb13307a71acd8fe81861d94ad54ab689df773318809eed3cbe794b4492dae4
+size 11422654

tokenizer_config.json

@@ -0,0 +1,240 @@
+{
+  "add_bos_token": false,
+  "add_prefix_space": false,
+  "added_tokens_decoder": {
+    "151643": {
+      "content": "<|endoftext|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151644": {
+      "content": "<|im_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151645": {
+      "content": "<|im_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151646": {
+      "content": "<|object_ref_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151647": {
+      "content": "<|object_ref_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151648": {
+      "content": "<|box_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151649": {
+      "content": "<|box_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151650": {
+      "content": "<|quad_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151651": {
+      "content": "<|quad_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151652": {
+      "content": "<|vision_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151653": {
+      "content": "<|vision_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151654": {
+      "content": "<|vision_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151655": {
+      "content": "<|image_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151656": {
+      "content": "<|video_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151657": {
+      "content": "<tool_call>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151658": {
+      "content": "</tool_call>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151659": {
+      "content": "<|fim_prefix|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151660": {
+      "content": "<|fim_middle|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151661": {
+      "content": "<|fim_suffix|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151662": {
+      "content": "<|fim_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151663": {
+      "content": "<|repo_name|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151664": {
+      "content": "<|file_sep|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151665": {
+      "content": "<tool_response>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151666": {
+      "content": "</tool_response>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151667": {
+      "content": "<think>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "151668": {
+      "content": "</think>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    }
+  },
+  "additional_special_tokens": [
+    "<|im_start|>",
+    "<|im_end|>",
+    "<|object_ref_start|>",
+    "<|object_ref_end|>",
+    "<|box_start|>",
+    "<|box_end|>",
+    "<|quad_start|>",
+    "<|quad_end|>",
+    "<|vision_start|>",
+    "<|vision_end|>",
+    "<|vision_pad|>",
+    "<|image_pad|>",
+    "<|video_pad|>"
+  ],
+  "bos_token": null,
+  "chat_template": "{%- if tools %}\n    {{- '<|im_start|>system\\n' }}\n    {%- if messages[0]['role'] == 'system' %}\n        {{- messages[0]['content'] }}\n    {%- else %}\n        {{- '' }}\n    {%- endif %}\n    {{- \"\\n\\n# Tools\\n\\nYou may call one or more functions to assist with the user query.\\n\\nYou are provided with function signatures within <tools></tools> XML tags:\\n<tools>\" }}\n    {%- for tool in tools %}\n        {{- \"\\n\" }}\n        {{- tool | tojson }}\n    {%- endfor %}\n    {{- \"\\n</tools>\\n\\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\\n<tool_call>\\n{\\\"name\\\": <function-name>, \\\"arguments\\\": <args-json-object>}\\n</tool_call><|im_end|>\\n\" }}\n{%- else %}\n    {%- if messages[0]['role'] == 'system' %}\n        {{- '<|im_start|>system\\n' + messages[0]['content'] + '<|im_end|>\\n' }}\n  {%- endif %}\n{%- endif %}\n{%- for message in messages %}\n    {%- if (message.role == \"user\") or (message.role == \"system\" and not loop.first) %}\n        {{- '<|im_start|>' + message.role + '\\n' + message.content + '<|im_end|>' + '\\n' }}\n    {%- elif message.role == \"assistant\" and not message.tool_calls %}\n        {%- set content = message.content %}\n        {%- if not loop.last %}\n            {%- set content = message.content.split('</think>')[-1].lstrip('\\n') %}\n        {%- endif %}\n        {{- '<|im_start|>' + message.role + '\\n' + content + '<|im_end|>' + '\\n' }}\n    {%- elif message.role == \"assistant\" %}\n        {%- set content = message.content %}\n        {%- if not loop.last %}\n            {%- set content = message.content.split('</think>')[-1].lstrip('\\n') %}\n        {%- endif %}\n        {{- '<|im_start|>' + message.role }}\n        {%- if message.content %}\n            {{- '\\n' + content }}\n        {%- endif %}\n        {%- for tool_call in message.tool_calls %}\n            {%- if tool_call.function is defined %}\n                {%- set tool_call = tool_call.function %}\n            {%- endif %}\n            {{- '\\n<tool_call>\\n{\"name\": \"' }}\n            {{- tool_call.name }}\n            {{- '\", \"arguments\": ' }}\n            {{- tool_call.arguments | tojson }}\n            {{- '}\\n</tool_call>' }}\n        {%- endfor %}\n        {{- '<|im_end|>\\n' }}\n    {%- elif message.role == \"tool\" %}\n        {%- if (loop.index0 == 0) or (messages[loop.index0 - 1].role != \"tool\") %}\n            {{- '<|im_start|>user' }}\n        {%- endif %}\n        {{- '\\n<tool_response>\\n' }}\n        {{- message.content }}\n        {{- '\\n</tool_response>' }}\n        {%- if loop.last or (messages[loop.index0 + 1].role != \"tool\") %}\n            {{- '<|im_end|>\\n' }}\n        {%- endif %}\n    {%- endif %}\n{%- endfor %}\n{%- if add_generation_prompt %}\n    {{- '<|im_start|>assistant\\n' }}\n{%- endif %}\n",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|im_end|>",
+  "errors": "replace",
+  "extra_special_tokens": {},
+  "model_max_length": 131072,
+  "pad_token": "<|endoftext|>",
+  "split_special_tokens": false,
+  "tokenizer_class": "Qwen2Tokenizer",
+  "unk_token": null
+}