import faiss 

# 建立索引, 定义为dimension d = 128
index = faiss.IndexFlatL2(d)

 # add vectors, xb 为 (100000,128)大小的numpy
index.add(xb)                 
print(index.ntotal) 
# 索引中向量的数量, 输出100000

# 求4-近邻
k = 4
# xq为query embedding, 大小为(10000,128)
D, I = index.search(xq, k)     
## D shape (10000,4)，表示每个返回点的embedding 与 query embedding的距离,
## I shape (10000,4)，表示和query embedding最接近的k个物品id，
print(I[:5])

import faiss
nlist = 100
k = 4

# 建立索引, 定义为dimension d = 128
quantizer = faiss.IndexFlatL2(d)

# 使用欧式距离 L2 建立索引。
index = faiss.IndexIVFFlat(quantizer, d, nlist, faiss.METRIC_L2)

## xb: (100000,128)
index.train(xb) 
index.add(xb)                
index.nprobe = 10  # 默认 nprobe 是 1 ,可以设置的大一些试试
D, I = index.search(xq, k)
print(I[-5:])   # 最后五次查询的结果

import faiss

nlist = 100
# 每个向量分8段
m = 8 
# 求4-近邻
k = 4 
quantizer = faiss.IndexFlatL2(d)    # 内部的索引方式依然不变
index = faiss.IndexIVFPQ(quantizer, d, nlist, m, 8) # 每个向量都被编码为8个字节大小
index.train(xb)
index.add(xb)
index.nprobe = 10                
D, I = index.search(xq, k)  # 检索
print(I[-5:])

def load_dataset(part='test') -> List[str]:
    ds = datasets.load_dataset("ag_news")
    list_str = [r['text'] for r in ds[part]]
    return list_str
    
list_str = load_dataset(part='train')
print(f'{len(list_str)}')
for s in list_str[:3]:
    print(s)
    print('\n')

index = load_index('news_train.index')
list_id = query(model, index, 'how to make money')
for id in list_id:
    print(list_str[id])

def train_flat(index_name, id_list, embedding_list, num_clusters):
    import numpy as np
    import faiss

    dim = 768
    m = 16
    
    embeddings = np.asarray(embedding_list)
    
    quantiser = faiss.IndexFlatIP(dim)
    index = faiss.IndexIVFFlat(quantiser, dim, num_clusters, faiss.METRIC_INNER_PRODUCT)
    index.train(embeddings)  ## clustering
    
    ids = np.arange(0, len(id_list))
    ids = np.asarray(ids.astype('int64'))
    
    index.add_with_ids(embeddings, ids)
    print(index.is_trained) 
    print("Total Number of Embeddings in the index", index.ntotal)
    faiss.write_index(index, index_name)

def query(model, index, query_str: str) -> List[int]:
    topk = 5
    q_embed = model.encode([query_str])
    D, I = index.search(q_embed, topk)
    print(D)
    print(I)
    return I[0].tolist()

from transformers import BertTokenizer, BertModel

model_name = 'hfl/chinese-bert-wwm'
tokenizer = BertTokenizer.from_pretrained(model_name)
model = BertModel.from_pretrained(model_name)

from sentence_transformers import SentenceTransformer
model = SentenceTransformer('paraphrase-MiniLM-L6-v2')

#Our sentences we like to encode
sentences = ['This framework generates embeddings for each input sentence',
    'Sentences are passed as a list of string.',
    'The quick brown fox jumps over the lazy dog.']

#Sentences are encoded by calling model.encode()
embeddings = model.encode(sentences)

#Print the embeddings
for sentence, embedding in zip(sentences, embeddings):
    print("Sentence:", sentence)
    print("Embedding:", embedding)
    print("")

from transformers import AutoTokenizer, AutoModel
import torch

#Mean Pooling - Take attention mask into account for correct averaging
def mean_pooling(model_output, attention_mask):
    token_embeddings = model_output[0] #First element of model_output contains all token embeddings
    input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
    sum_embeddings = torch.sum(token_embeddings * input_mask_expanded, 1)
    sum_mask = torch.clamp(input_mask_expanded.sum(1), min=1e-9)
    return sum_embeddings / sum_mask

#Sentences we want sentence embeddings for
sentences = ['This framework generates embeddings for each input sentence',
             'Sentences are passed as a list of string.',
             'The quick brown fox jumps over the lazy dog.']

#Load AutoModel from huggingface model repository
tokenizer = AutoTokenizer.from_pretrained("sentence-transformers/all-MiniLM-L6-v2")
model = AutoModel.from_pretrained("sentence-transformers/all-MiniLM-L6-v2")

#Tokenize sentences
encoded_input = tokenizer(sentences, padding=True, truncation=True, max_length=128, return_tensors='pt')

#Compute token embeddings
with torch.no_grad():
    model_output = model(**encoded_input)

#Perform pooling. In this case, mean pooling
sentence_embeddings = mean_pooling(model_output, encoded_input['attention_mask'])

sentences = [
    '今天晚上想吃牛排',
    'MyEncyclopedia公众号全栈人工智能',
    '人工智能需要懂很多数学么',
    '上海疫情有完没完',
    '教育部：连续7天社会面无疫情 高校可组织校园招聘',
    '福建舰"下水！100秒看中国航母高光时刻',
    '医保承担多少核酸检测费用？压力多大？',
    '张家口过度防疫整改后又被曝光：要证明牛是阴性',
    '上海多家银行天天排队爆满 有老人凌晨2点开始排队',
    'A股不惧海外暴跌！走出独立行情沪指收复3300点',
    '俄方称已准备好重启俄乌和谈',
    '《自然》：奥密克戎感染后嗅觉丧失症状比原来少了'
]

q1 = '码农的春天来了么'
q2 = '国际局势'
q3 = '健康'

from sentence_transformers import SentenceTransformer

model_name = 'hfl/chinese-bert-wwm'
model = SentenceTransformer(model_name)

sentences = [
    '今天晚上想吃牛排',
    'MyEncyclopedia公众号全栈人工智能',
    '人工智能需要懂很多数学么',
    '上海疫情有完没完',
    '教育部：连续7天社会面无疫情 高校可组织校园招聘',
    '福建舰"下水！100秒看中国航母高光时刻',
    '医保承担多少核酸检测费用？压力多大？',
    '张家口过度防疫整改后又被曝光：要证明牛是阴性',
    '上海多家银行天天排队爆满 有老人凌晨2点开始排队',
    'A股不惧海外暴跌！走出独立行情沪指收复3300点',
    '俄方称已准备好重启俄乌和谈',
    '《自然》：奥密克戎感染后嗅觉丧失症状比原来少了'
]
sentence_embeddings = model.encode(sentences)

q1 = '码农的春天来了么'
q2 = '国际局势'
q3 = '健康'
queries = [q1, q2, q3]
query_embeddings = model.encode(queries)

import scipy

number_top_matches = 3
for query, query_embedding in zip(queries, query_embeddings):
    distances = scipy.spatial.distance.cdist([query_embedding], sentence_embeddings, "cosine")[0]
    results = zip(range(len(distances)), distances)
    results = sorted(results, key=lambda x: x[1])
    print("\nQuery:", query)
    print("\nTop {} most similar sentences in corpus:".format(number_top_matches))

    for idx, distance in results[0:number_top_matches]:
        print(sentences[idx].strip(), "(Cosine Score: %.4f)" % (1-distance))

from flair.models import TARSClassifier
from flair.data import Sentence
import flair, torch
flair.device = torch.device('cpu')

text = 'Spain beat Swiss for first Nations League win'
tars = TARSClassifier.load('tars-base')
sentence = Sentence(text)
classes = ['politics', 'sports', 'health']
tars.predict_zero_shot(sentence, classes)

print(sentence)
print(sentence.to_dict())

yahoo = load_dataset('yahoo_answers_topics')

[
'Society & Culture', 
'Science & Mathematics', 
'Health', 
'Education & Reference', 
'Computers & Internet', 
'Sports', 
'Business & Finance', 
'Entertainment & Music', 
'Family & Relationships', 
'Politics & Government'
]

def evaluate_flair(dataset, default_name='neutral'):
    classifier = TARSClassifier.load('tars-base')
    total, correct = 0, 0
    for item, gold_label_idx in tqdm(zip(dataset["test_texts"], dataset["test_labels"]),
                                     total=len(dataset["test_texts"])):
        sentence = Sentence(item)
        classifier.predict_zero_shot(sentence, dataset["class_names"])
        sorted_labels = sorted(sentence.to_dict()['all labels'], key=lambda k: k['confidence'], reverse=True)
        gold_label = dataset["class_names"][gold_label_idx]
        if len(sorted_labels) > 0:
            predicted_label = sorted_labels[0]['value']
        else:
            predicted_label = default_name
        if predicted_label == gold_label:
            correct += 1
        total += 1

    return correct / total
    

def evaluate_huggingface(dataset):
    classifier = pipeline("zero-shot-classification", device=-1)
    correct = 0
    predictions, gold_labels = [], []
    for text, gold_label_idx in tqdm(zip(dataset["test_texts"], dataset["test_labels"]),
                                     total=len(dataset["test_texts"])):

        result = classifier(text, dataset["class_names"], multi_label=False)
        predicted_label = result['labels'][0]

        gold_label = dataset["class_names"][gold_label_idx]

        predictions.append(predicted_label)
        gold_labels.append(gold_label)

        if predicted_label == gold_label:
            correct += 1

    accuracy = correct / len(predictions)
    return accuracy