%load_ext autoreload
%autoreload 2
import mycode.vap as vap

# The goal is to show how to use the hex embedding and name some pitfalls

from vespa.package import (ApplicationPackage, Field, Schema, Document)
ap = ApplicationPackage(
    name="hex",
    schema=[
        Schema(
            name="hex",
            document=Document(
                fields=[
                    Field(
                        name="embedding_float",
                        type="tensor<float>(d0[1])",
                        indexing=["attribute", "summary"],
                    ),
                    Field(
                        name="embedding_bfloat16",
                        type="tensor<bfloat16>(d0[1])",
                        indexing=["attribute", "summary"],
                    )
                ]
            )
        )
    ]
)

from vespa.deployment import VespaDocker

# In case running colima on macos run the following
# !sudo ln -sf $HOME/.colima/default/docker.sock /var/run/docker.sock
vespa_docker = VespaDocker(
    container_image="vespaengine/vespa:8.625.17",
)

vespa_client = vespa_docker.deploy(application_package=ap)

import struct

def float_to_hex(f: float) -> str:
    return format(struct.unpack('=I', struct.pack('=f', f))[0], '08X')

def hex_to_float(hex_str: str) -> float:
    i = int(hex_str, 16)
    return struct.unpack('=f', struct.pack('=I', i))[0]

def float_to_bf16_hex(f: float) -> str:
    f32_bits = struct.unpack('=I', struct.pack('=f', f))[0]
    bf16_bits = f32_bits >> 16
    return format(bf16_bits, '04X')

def bf16_hex_to_float(hex_str: str) -> float:
    bf16_bits = int(hex_str, 16)
    f32_bits = bf16_bits << 16
    return struct.unpack('=f', struct.pack('=I', f32_bits))[0]

num = 1.23456
print(f'Original_float={num}, hex={float_to_hex(num)}, decoded_float={hex_to_float(float_to_hex(num))}')
print(f'Original_float={num}, hex={float_to_bf16_hex(num)}, decoded_float={bf16_hex_to_float(float_to_bf16_hex(num))}')

# Feed the same float number to both fields
vespa_client.feed_iterable([
    {
        'id': '1',
        'fields': {
            'embedding_float': [num],
            'embedding_bfloat16': [num],
        }
    }
], schema="hex", namespace="hex", callback=vap.feed_callback)

vespa_client.get_data(data_id='1', schema="hex", namespace="hex").json

vespa_client.feed_iterable([
    {
        'id': '2',
        'fields': {
            'embedding_float': float_to_hex(num),
        }
    }
], schema="hex", namespace="hex", callback=vap.feed_callback)

vespa_client.get_data(data_id='2', schema="hex", namespace="hex").json

# now let's try feeding hex encoded float into a bfloat16 field
vespa_client.feed_iterable([
    {
        'id': '3',
        'fields': {
            'embedding_bfloat16': float_to_hex(num),
        }
    }
], schema="hex", namespace="hex", callback=vap.feed_callback)

# We've got a strange out of bound exception
# Fetch the doc anyway:
vespa_client.get_data(data_id='3', schema="hex", namespace="hex").json

# The doc is not present, as expected because Vespa threw an exception
# now let's try feeding hex encoded float into a bfloat16 field
vespa_client.feed_iterable([
    {
        'id': '4',
        'fields': {
            'embedding_bfloat16': float_to_bf16_hex(num),
        }
    }
], schema="hex", namespace="hex", callback=vap.feed_callback)

vespa_client.get_data(data_id='4', schema="hex", namespace="hex").json

# As expected.