jax jaxlib flax transformers streamlit tqdm nmslib matplotlib pandas requests