# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

import torch, datasets, transformers
import numpy as np
dataset = datasets.load_from_disk("clip_feature_featonly")
print(dataset)

from sklearn.cluster import SpectralClustering, KMeans
import os.path as osp
datadict = dataset.to_dict()

from itertools import islice
print(datadict.keys())
print(datadict["feature"][0])

embedded: np.ndarray = np.stack(datadict["feature"])
print(embedded.shape, embedded)

from sklearn.neighbors import NearestNeighbors
nn = NearestNeighbors(n_neighbors=10)
nn.fit(embedded)

knns = nn.kneighbors(embedded[0:1], 5)
print(knns)

datadict["feature"]=embedded

datadict["feature"].shape

np.save("clip_feature_for_knn.npy", datadict)

import pickle
with open("clip_feature_for_knn.pkl", "wb") as f:
    pickle.dump(datadict, f)



StarcoderdataJupyterScriptsDedupFiltered