from transformers import AutoTokenizer, AutoModel
import torch.nn as nn

tokenizer = AutoTokenizer.from_pretrained("sap-ai-research/miCSE")
model = AutoModel.from_pretrained("sap-ai-research/miCSE")

max_length = 32
sentences = [
    "This is a sentence for testing miCSE.", 
    "This is yet another test sentence for the mutual information Contrastive Sentence Embeddings model."
]

batch = tokenizer.batch_encode_plus(
            sentences,
            return_tensors='pt',
            padding=True,
            max_length=max_length,
            truncation=True
        )

with torch.no_grad():
    outputs = model(**batch, output_hidden_states=True, return_dict=True)
embeddings = outputs.last_hidden_state[:,0]

sim = nn.CosineSimilarity(dim=-1)
cos_sim = sim(embeddings.unsqueeze(1), embeddings.unsqueeze(0))
print(f"相似度: {cos_sim[0,1].detach().item()}")

from transformers import AutoTokenizer, AutoModel
import torch.nn as nn
import torch
import numpy as np
import tqdm
from datasets import load_dataset
import umap
import umap.plot as umap_plot

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
tokenizer = AutoTokenizer.from_pretrained("/Users/d065243/miCSE")
model = AutoModel.from_pretrained("/Users/d065243/miCSE").to(device)

dataset = load_dataset("tweet_eval", "sentiment")
batch_size = 50
max_length = 128

embedding_stack = []
classes = []
for i in tqdm.notebook.tqdm(range(0,len(dataset['train']),batch_size)):
    batch = tokenizer.batch_encode_plus(
                dataset['train'][i:i+batch_size]['text'],
                return_tensors='pt',
                padding=True,
                max_length=max_length,
                truncation=True
            ).to(device)
    classes += dataset['train'][i:i+batch_size]['label']
    
    with torch.no_grad():
        outputs = model(**batch, output_hidden_states=True, return_dict=True)
        embedding_stack.append(outputs.last_hidden_state[:,0].cpu())

embeddings = torch.vstack(embedding_stack)
umap_model = umap.UMAP(n_neighbors=250, n_components=2, metric='cosine').fit(embeddings)
umap_plot.points(umap_model, labels=np.array(classes), theme='fire')

from sentence_transformers import SentenceTransformer, models
import torch.nn as nn

model = SentenceTransformer(modules=[
    models.Transformer('sap-ai-research/miCSE', max_seq_length=32),
    models.Pooling(models.Transformer.get_word_embedding_dimension())
])

sentences1 = ["测试miCSE的句子", "使用互信息对比句子嵌入模型"]
sentences2 = ["miCSE测试", "与miCSE的相似度"]

embeddings1 = model.encode(sentences1, convert_to_tensor=True)
embeddings2 = model.encode(sentences2, convert_to_tensor=True)

cos_sim = nn.CosineSimilarity(dim=-1)(embeddings1, embeddings2)
for i in range(len(sentences1)):
    print(f"相似度{cos_sim[i]:.2f}: {sentences1[i]} <<对比>> {sentences2[i]}")

+-------+-------+-------+-------+-------+--------------+-----------------+--------+
| STS12 | STS13 | STS14 | STS15 | STS16 | STSBenchmark | SICKRelatedness | S.Avg. |
+-------+-------+-------+-------+-------+--------------+-----------------+--------+
| 71.71 | 83.09 | 75.46 | 83.13 | 80.22 |    79.70     |      73.62      | 78.13  |
+-------+-------+-------+-------+-------+--------------+-----------------+--------+