import streamlit as st
from huggingface_hub import HfFileSystem
from flax.serialization import msgpack_restore, from_state_dict

fs = HfFileSystem()

with fs.open("PrakhAI/HelloWorld/checkpoint.msgpack", "rb") as f:
    state = msgpack_restore(f.read())

print(type(state))
print(state)
print(state.keys)
print(state['params'])
print(state['opt_state'])

# logits = state.apply_fn({'params': state.params}, batch['image'])
# print(logits)

x = st.slider('Select a value')
st.write(x, 'squared is', x * x)