# -- Code Cell --
import pandas as pd
import torch
from transformers import AutoTokenizer, AutoModel
import numpy as np
import os

#Load data and model
test_df = pd.read_csv("test_data.csv")

model_path = os.path.expanduser("/opt/hfmodels/BAAI--bge-small-en-v1.5")

model = AutoModel.from_pretrained("BAAI/bge-small-en-v1.5")
tokenizer = AutoTokenizer.from_pretrained("BAAI/bge-small-en-v1.5")

#Example of using model to get sentence embeddings

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)
model.eval()

sample_text = test_df['text'].iloc[0]

with torch.no_grad():
    encoded_input = tokenizer(sample_text, return_tensors="pt").to(device)
    model_output = model(**encoded_input)

# [CLS] token at index 0 stores embedding of the whole sentence
embedding = model_output[0][:, 0]

# -- Code Cell --
embeddings = []
for i in range(len(test_df)):
    with torch.no_grad():
        text = test_df['text'][i]
        encoded_input = tokenizer(text, return_tensors="pt").to(device)
        model_output = model(**encoded_input)
        embedding = model_output[0][:, 0]
        embeddings.append(embedding.cpu())

# -- Code Cell --
len(test_df)

# -- Code Cell --
len(embeddings)

# -- Code Cell --
embeddings[1]

# -- Code Cell --
embeddings_squee = []
for i in range(len(embeddings)):
    embeddings_squee.append(embeddings[i].squeeze())

# -- Code Cell --
from sklearn.decomposition import PCA
pca = PCA(n_components=50)
X_pca = pca.fit_transform(embeddings_squee)
ceva = pca.inverse_transform(X_pca)

# -- Code Cell --
ceva

# -- Code Cell --
error = np.mean((embeddings_squee - ceva) ** 2, axis=1)

# -- Code Cell --
import umap
from sklearn.cluster import HDBSCAN
umapp = umap.UMAP(n_components=10, n_neighbors=30, min_dist=0.0, metric='cosine', random_state=42)
X = np.array(embeddings_squee)
X_umap = umapp.fit_transform(X)
clustere = HDBSCAN(min_cluster_size=100, min_samples=10, metric='euclidean')
hdb_labels = clustere.fit_predict(X_umap)

# -- Code Cell --
import matplotlib.pyplot as plt
X_vis = umap.UMAP(n_components=2, n_neighbors=30, min_dist=0.1, metric='cosine', random_state=42).fit_transform(X)
plt.figure(figsize=(12, 8))
plt.scatter(X_vis[:, 0], X_vis[:, 1], c=hdb_labels)
plt.show()

# -- Code Cell --
labels = [1 if hdb_labels[i] == 0 else 0 for i in range(len(hdb_labels))]

# -- Code Cell --
thre = 0.1
outliters = error>thre
labels = [1 if outliters[i] else 0 for i in range(len(embeddings_squee))]

# -- Code Cell --
labels

# -- Code Cell --
sub1 = pd.DataFrame({
    "subtaskID":1,
    "datapointID":test_df['datapointID'],
    "answer":labels
})
sub2 = pd.DataFrame({
    "subtaskID":2,
    "datapointID":test_df['datapointID'],
    "answer":hdb_labels
})
final = pd.concat([sub1,sub2]).to_csv("subs.csv",index=False)

# -- Code Cell --
