import gradio as gr
import spaces
import torch
import vdf_io

zero = torch.Tensor([0]).cuda()
print(zero.device)  # <-- 'cpu' 🤔

print(vdf_io.__version__)


@spaces.GPU
def greet(n):
    print(zero.device)  # <-- 'cuda:0' 🤗
    return f"Hello {zero + n} Tensor"


def reembed_dataset():
    import datasets
    # model 
    # embeddings = model.embed(ds)
    # new_embeddings = model.reembed(embeddings)

    # datasets.save_dataset(new_embeddings)


def reembed_main():
    download_dataset()
    reembed_dataset()


def download_dataset():
    import datasets

    # ds = datasets.load_dataset()


demo = gr.Interface(
    fn=reembed_main,
    inputs=[
        # dataset name
        gr.inputs.Textbox(label="Dataset name"),
        # embedding model
        gr.inputs.Textbox(label="Embedding model"),
        # output username
        gr.inputs.Textbox(label="Output username"),
    ],
    outputs=gr.outputs.Textbox(label="Output"),
    title="Re-Embedder",
    description="Re-embed a dataset using a given model and output to a new username's account",
)
demo.launch()