Skip to content

Instantly share code, notes, and snippets.

@szmeku
Created July 11, 2024 22:04
Show Gist options
  • Save szmeku/82475128388e90d434805c1c06eee35c to your computer and use it in GitHub Desktop.
Save szmeku/82475128388e90d434805c1c06eee35c to your computer and use it in GitHub Desktop.

Sure, for running an LLM locally on an iOS device, we can use a smaller, more efficient model like GPT-2 or a distilled version. For simplicity, I will demonstrate using the CoreML framework with a pre-trained GPT-2 model that has been converted to CoreML format. This example will include the necessary steps to integrate the model and measure the response time.

Complete Example for iOS App with Swift, Metal, and a Local GPT-2 Model

1. Convert a Pre-trained GPT-2 Model to CoreML

First, you need to convert a pre-trained GPT-2 model to CoreML format. This is typically done outside Xcode using a Python script. Here is a basic script to convert a GPT-2 model to CoreML:

# Install the required packages
# pip install transformers coremltools

import coremltools as ct
from transformers import GPT2LMHeadModel, GPT2Tokenizer

model_name = "gpt2"
tokenizer = GPT2Tokenizer.from_pretrained(model_name)
model = GPT2LMHeadModel.from_pretrained(model_name)

# Trace the model with a dummy input
input_text = "Convert this text"
input_ids = tokenizer.encode(input_text, return_tensors="pt")

traced_model = torch.jit.trace(model, input_ids)
coreml_model = ct.convert(
    traced_model,
    inputs=[ct.TensorType(name="input_ids", shape=input_ids.shape)],
)

coreml_model.save("GPT2.mlmodel")

Ensure you have the GPT2.mlmodel file ready.

2. Add the Model to Your Xcode Project

Drag and drop the GPT2.mlmodel file into your Xcode project.

3. Replace ContentView.swift

Replace the contents of ContentView.swift with the following code:

import SwiftUI

struct ContentView: View {
    var body: some View {
        Text("LLM Response Time Test")
            .padding()
            .onAppear {
                LLMManager.shared.runLLMTest()
            }
    }
}

struct ContentView_Previews: PreviewProvider {
    static var previews: some View {
        ContentView()
    }
}

4. Create LLMManager.swift

Create a new file named LLMManager.swift and add the following code:

import Foundation
import CoreML

class LLMManager {
    static let shared = LLMManager()
    
    private init() {}
    
    func runLLMTest() {
        let input = "Fixed prompt for the language model"
        
        // Measure the time of LLM response
        let startTime = CFAbsoluteTimeGetCurrent()
        
        // Perform LLM inference
        let output = processLLM(input: input)
        
        let endTime = CFAbsoluteTimeGetCurrent()
        let responseTime = endTime - startTime
        
        print("LLM Response: \(output)")
        print("LLM Response Time: \(responseTime) seconds")
    }
    
    private func processLLM(input: String) -> String {
        guard let model = try? GPT2(configuration: MLModelConfiguration()) else {
            return "Failed to load model"
        }
        
        guard let input_ids = try? GPT2Input(text: input) else {
            return "Failed to create input"
        }
        
        guard let output = try? model.prediction(input: input_ids) else {
            return "Failed to make prediction"
        }
        
        return output.text
    }
}

5. Create MetalRenderer.swift

Create a new file named MetalRenderer.swift and add the following code:

import Metal
import MetalKit

class MetalRenderer: NSObject, MTKViewDelegate {
    var device: MTLDevice!
    var commandQueue: MTLCommandQueue!
    
    init(mtkView: MTKView) {
        self.device = MTLCreateSystemDefaultDevice()
        mtkView.device = self.device
        mtkView.colorPixelFormat = .bgra8Unorm
        self.commandQueue = device.makeCommandQueue()
        super.init()
        mtkView.delegate = self
    }
    
    func draw(in view: MTKView) {
        guard let drawable = view.currentDrawable,
              let descriptor = view.currentRenderPassDescriptor else { return }
        
        let commandBuffer = commandQueue.makeCommandBuffer()
        let renderEncoder = commandBuffer?.makeRenderCommandEncoder(descriptor: descriptor)
        
        renderEncoder?.endEncoding()
        commandBuffer?.present(drawable)
        commandBuffer?.commit()
    }
    
    func mtkView(_ view: MTKView, drawableSizeWillChange size: CGSize) {}
}

6. Integrate Metal into SwiftUI

Edit the main app file, which should be named something like YourProjectNameApp.swift, to set up the Metal view.

import SwiftUI

@main
struct YourProjectNameApp: App {
    var body: some Scene {
        WindowGroup {
            MetalView()
        }
    }
}

struct MetalView: UIViewRepresentable {
    func makeCoordinator() -> Coordinator {
        Coordinator(self)
    }
    
    func makeUIView(context: Context) -> MTKView {
        let mtkView = MTKView()
        context.coordinator.renderer = MetalRenderer(mtkView: mtkView)
        return mtkView
    }
    
    func updateUIView(_ uiView: MTKView, context: Context) {}
    
    class Coordinator: NSObject {
        var parent: MetalView
        var renderer: MetalRenderer?
        
        init(_ parent: MetalView) {
            self.parent = parent
        }
    }
}

Explanation

  1. ContentView.swift: This is the main view that displays a simple text and triggers the LLM test on appearance.
  2. LLMManager.swift: This file contains a singleton LLMManager class that runs a fixed prompt through the CoreML model and measures the response time.
  3. MetalRenderer.swift: This file sets up a basic Metal renderer that clears the screen. The Metal rendering is minimal since the focus is on the LLM response.
  4. YourProjectNameApp.swift: This file integrates the Metal view into the SwiftUI app.

This example uses a CoreML model converted from GPT-2. The CoreML model allows you to perform inference locally on the device without requiring an external API. The provided code demonstrates how to integrate this model into an iOS app, measure the response time, and perform basic Metal rendering.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment