v2: Initial clean branch with unit tests + CI/CD pipeline
- Started from ac75faa (initial E4B-MarkBase integration)
- Kept Sources/ (all engine code) + Package.swift + .gitignore
- Removed all ad-hoc tests, documentation, scripts, Python files
- Added Tests/00_Unit/ (MathTest, TokenizerTest, SamplerTest)
- Added .gitea/workflows/ci.yaml (build + unit tests + lint)
- Added Scripts/check_resources.sh (memory-aware test runner)
- Added Tests/Manifest.json (resource requirements for all tests)
- Focus: 4-bit quantized models only
This commit is contained in:
@@ -0,0 +1,34 @@
|
||||
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) }
|
||||
}
|
||||
Reference in New Issue
Block a user