try embed

app.py

@@ -2,6 +2,8 @@ import gradio as gr
 import spaces
 import torch
 import vdf_io
+from sentence_transformers import SentenceTransformer
+from rich import print as rprint
 
 zero = torch.Tensor([0]).cuda()
 print(zero.device)  # <-- 'cpu' 🤔
@@ -15,24 +17,28 @@ def greet(n):
     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)
+@spaces.GPU
+def reembed_dataset(ds, model):
+    model = SentenceTransformer(model, device=zero.device)
+    rprint(model)
+    rprint(model.encode("Hello, World!"))
+    ds.map(lambda x: model.encode(x["text"]))
+    rprint(ds[0])
 
 
-def reembed_main():
-    download_dataset()
-    reembed_dataset()
+def reembed_main(dataset_name, embedding_model, output_username):
+    print(f"{dataset_name=}, {embedding_model=}, {output_username=}")
+    ds = download_dataset(dataset_name)
+    reembed_dataset(ds, model=embedding_model)
+    return "Dataset re-embedded successfully"
 
 
-def download_dataset():
+def download_dataset(dataset_name):
     import datasets
 
-    # ds = datasets.load_dataset()
+    ds = datasets.load_dataset(dataset_name)
+    print(len(ds))
+    return ds
 
 
 demo = gr.Interface(