1 /*
 2  * Copyright (c) 2025, Oracle and/or its affiliates. All rights reserved.
 3  * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
 4  *
 5  * This code is free software; you can redistribute it and/or modify it
 6  * under the terms of the GNU General Public License version 2 only, as
 7  * published by the Free Software Foundation.
 8  *
 9  * This code is distributed in the hope that it will be useful, but WITHOUT
10  * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
11  * FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License
12  * version 2 for more details (a copy is included in the LICENSE file that
13  * accompanied this code).
14  *
15  * You should have received a copy of the GNU General Public License version
16  * 2 along with this work; if not, write to the Free Software Foundation,
17  * Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301 USA.
18  *
19  * Please contact Oracle, 500 Oracle Parkway, Redwood Shores, CA 94065 USA
20  * or visit www.oracle.com if you need additional information or have any
21  * questions.
22  */
23 package oracle.code.onnx.genai;
24 
25 import java.io.IOException;
26 import java.io.RandomAccessFile;
27 import java.lang.foreign.Arena;
28 import java.lang.foreign.MemorySegment;
29 import java.nio.channels.FileChannel;
30 import java.util.stream.LongStream;
31 import oracle.code.onnx.Tensor;
32 
33 public final class TensorDataStream {
34 
35 
36     private final Arena arena;
37     private final MemorySegment data;
38     private long offset;
39 
40     public TensorDataStream(Arena arena, String dataFilePath) throws IOException {
41         this.arena = arena;
42         try (var dataFile = new RandomAccessFile(dataFilePath, "r")) {
43             this.data = dataFile.getChannel().map(FileChannel.MapMode.READ_ONLY, 0, dataFile.length(), arena);
44         }
45     }
46 
47     public TensorDataStream(Arena arena, MemorySegment data) {
48         this.arena = arena;
49         this.data = data;
50     }
51 
52     public <T> Tensor<T> nextTensor(Tensor.ElementType type, long... shape) {
53         long size = type.bitSize() * LongStream.of(shape).reduce(1l, (a, b) -> a * b) / 8l;
54         Tensor<T> tensor = new Tensor<>(arena, data.asSlice(offset, size), type, shape);
55         offset += size;
56         return tensor;
57     }
58 }