the-algorithm/home-mixer/server/src/main/scala/com/twitter/home_mixer/util/TensorFlowUtil.scala
twitter-team ef4c5eb65e Twitter Recommendation Algorithm
Please note we have force-pushed a new initial commit in order to remove some publicly-available Twitter user information. Note that this process may be required in the future.
2023-03-31 17:36:31 -05:00

33 lines
806 B
Scala

package com.twitter.home_mixer.util
import com.twitter.ml.api.thriftscala.FloatTensor
import com.twitter.ml.api.util.BufferToIterators.RichFloatBuffer
import java.nio.ByteBuffer
import java.nio.ByteOrder
/**
* Contains functionality to transform data records and Tensors
*/
object TensorFlowUtil {
private def skipEmbeddingBBHeader(bb: ByteBuffer): ByteBuffer = {
val bb_copy = bb.duplicate()
bb_copy.getLong()
bb_copy
}
private def byteBufferToFloatIterator(
bb: ByteBuffer
): Iterator[Float] = {
bb.order(ByteOrder.LITTLE_ENDIAN).asFloatBuffer.iterator
}
def embeddingByteBufferToFloatTensor(
bb: ByteBuffer
): FloatTensor = {
val bb_content = skipEmbeddingBBHeader(bb)
FloatTensor(byteBufferToFloatIterator(bb_content).map(_.toDouble).toList)
}
}