<commit_before>import Foundation
import TensorSwift

public struct Classifier {
    public let W_conv1: Tensor
    public let b_conv1: Tensor
    public let W_conv2: Tensor
    public let b_conv2: Tensor
    public let W_fc1: Tensor
    public let b_fc1: Tensor
    public let W_fc2: Tensor
    public let b_fc2: Tensor
    
    public func classify(x_image: Tensor) -> Int {
        let h_conv1 = (x_image.conv2d(filter: W_conv1, strides: [1, 1, 1]) + b_conv1).relu
        let h_pool1 = h_conv1.maxPool(kernelSize: [2, 2, 1], strides: [2, 2, 1])
        
        let h_conv2 = (h_pool1.conv2d(filter: W_conv2, strides: [1, 1, 1]) + b_conv2).relu
        let h_pool2 = h_conv2.maxPool(kernelSize: [2, 2, 1], strides: [2, 2, 1])
        
        let h_pool2_flat = Tensor(shape: [1, Dimension(7*7*64)], elements: h_pool2.elements)
        let h_fc1 = (h_pool2_flat.matmul(W_fc1) + b_fc1).relu
        
        let y_conv = (h_fc1.matmul(W_fc2) + b_fc2).softmax

        return y_conv.elements.enumerate().maxElement { $0.1 < $1.1 }!.0
    }
}

extension Classifier {
    public init(path: String) {
        W_conv1 = Tensor(shape: [5, 5, 1, 32], elements: loadFloatArray(directory: path, file: "W_conv1"))
        b_conv1 = Tensor(shape: [32], elements: loadFloatArray(directory: path, file: "b_conv1"))
        W_conv2 = Tensor(shape: [5, 5, 32, 64], elements: loadFloatArray(directory: path, file: "W_conv2"))
        b_conv2 = Tensor(shape: [64], elements: loadFloatArray(directory: path, file: "b_conv2"))
        W_fc1 = Tensor(shape: [Dimension(7 * 7 * 64), 1024], elements: loadFloatArray(directory: path, file: "W_fc1"))
        b_fc1 = Tensor(shape: [1024], elements: loadFloatArray(directory: path, file: "b_fc1"))
        W_fc2 = Tensor(shape: [1024, 10], elements: loadFloatArray(directory: path, file: "W_fc2"))
        b_fc2 = Tensor(shape: [10], elements: loadFloatArray(directory: path, file: "b_fc2"))
    }
}

private func loadFloatArray(directory directory: String, file: String) -> [Float] {
    let data = NSData(contentsOfFile: directory.stringByAppendingPathComponent(file))!
    return Array(UnsafeBufferPointer(start: UnsafeMutablePointer<Float>(data.bytes), count: data.length / 4))
}<commit_msg>Update `Classifier` to use `Dimension`'s operator<commit_after>import Foundation
import TensorSwift

public struct Classifier {
    public let W_conv1: Tensor
    public let b_conv1: Tensor
    public let W_conv2: Tensor
    public let b_conv2: Tensor
    public let W_fc1: Tensor
    public let b_fc1: Tensor
    public let W_fc2: Tensor
    public let b_fc2: Tensor
    
    public func classify(x_image: Tensor) -> Int {
        let h_conv1 = (x_image.conv2d(filter: W_conv1, strides: [1, 1, 1]) + b_conv1).relu
        let h_pool1 = h_conv1.maxPool(kernelSize: [2, 2, 1], strides: [2, 2, 1])
        
        let h_conv2 = (h_pool1.conv2d(filter: W_conv2, strides: [1, 1, 1]) + b_conv2).relu
        let h_pool2 = h_conv2.maxPool(kernelSize: [2, 2, 1], strides: [2, 2, 1])
        
        let h_pool2_flat = Tensor(shape: [1, 7 * 7 * 64], elements: h_pool2.elements)
        let h_fc1 = (h_pool2_flat.matmul(W_fc1) + b_fc1).relu
        
        let y_conv = (h_fc1.matmul(W_fc2) + b_fc2).softmax

        return y_conv.elements.enumerate().maxElement { $0.1 < $1.1 }!.0
    }
}

extension Classifier {
    public init(path: String) {
        W_conv1 = Tensor(shape: [5, 5, 1, 32], elements: loadFloatArray(directory: path, file: "W_conv1"))
        b_conv1 = Tensor(shape: [32], elements: loadFloatArray(directory: path, file: "b_conv1"))
        W_conv2 = Tensor(shape: [5, 5, 32, 64], elements: loadFloatArray(directory: path, file: "W_conv2"))
        b_conv2 = Tensor(shape: [64], elements: loadFloatArray(directory: path, file: "b_conv2"))
        W_fc1 = Tensor(shape: [7 * 7 * 64, 1024], elements: loadFloatArray(directory: path, file: "W_fc1"))
        b_fc1 = Tensor(shape: [1024], elements: loadFloatArray(directory: path, file: "b_fc1"))
        W_fc2 = Tensor(shape: [1024, 10], elements: loadFloatArray(directory: path, file: "W_fc2"))
        b_fc2 = Tensor(shape: [10], elements: loadFloatArray(directory: path, file: "b_fc2"))
    }
}

private func loadFloatArray(directory directory: String, file: String) -> [Float] {
    let data = NSData(contentsOfFile: directory.stringByAppendingPathComponent(file))!
    return Array(UnsafeBufferPointer(start: UnsafeMutablePointer<Float>(data.bytes), count: data.length / 4))
}
{'fraction_non_alphanumeric': 0.13290004482294934, 'fraction_numerical': 0.049305244285073956, 'mean_word_length': 3.62008281573499, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 9, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2934120', 'n_tokens_mistral': 1676, 'n_tokens_neox': 1574, 'n_words': 405}