MarkBase Admin
ac75faa0cc
CI / build-and-test (push) Has been cancelled
- E4B-MarkBase model (42 layers, 4.4GB) loaded successfully - All Phase 1-6 tests passed (model loading, forward pass, vision/audio towers, token generation, performance) - All stress tests passed (5/5 in 127.6s) - Concurrent inference - Memory stress (67.5 tok/s, 0 NaN) - Continuous generation - Batch processing - Long-running stability - Swift Metal inference engine with multimodal support
35 lines
1.3 KiB
Swift
35 lines
1.3 KiB
Swift
import Foundation
|
|
|
|
/// Handles sharded SafeTensors models (with model.safetensors.index.json).
|
|
public final class SafeTensorsIndex {
|
|
public let weightMap: [String: String]
|
|
public let baseDir: String
|
|
|
|
/// Load the index file from a model directory.
|
|
public init(modelDir: String) throws {
|
|
let indexURL = URL(fileURLWithPath: modelDir).appendingPathComponent("model.safetensors.index.json")
|
|
let data = try Data(contentsOf: indexURL)
|
|
guard let json = try JSONSerialization.jsonObject(with: data) as? [String: Any],
|
|
let weightMap = json["weight_map"] as? [String: String]
|
|
else {
|
|
throw WeightError.invalidHeader("Index file missing weight_map")
|
|
}
|
|
self.weightMap = weightMap
|
|
self.baseDir = modelDir
|
|
}
|
|
|
|
/// All unique shard filenames referenced by the index.
|
|
public var shardFiles: Set<String> {
|
|
Set(weightMap.values)
|
|
}
|
|
|
|
/// Resolve a tensor name to its shard file path.
|
|
public func shardPath(for tensor: String) -> String? {
|
|
guard let shard = weightMap[tensor] else { return nil }
|
|
return (baseDir as NSString).appendingPathComponent(shard)
|
|
}
|
|
|
|
/// List all tensor names.
|
|
public var allTensors: [String] { Array(weightMap.keys) }
|
|
}
|