import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import numpy as np

class CrossEncoderReranker:
    def __init__(self, model_name='cross-encoder/ms-marco-MiniLM-L-6-v2'):
        """Cross-Encoderリランカーの初期化"""
        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
        self.model = AutoModelForSequenceClassification.from_pretrained(model_name)
        self.model.eval()

    def score(self, query, documents):
        """クエリと文書ペアのスコアを計算"""
        pairs = [[query, doc] for doc in documents]

        with torch.no_grad():
            inputs = self.tokenizer(
                pairs,
                padding=True,
                truncation=True,
                max_length=512,
                return_tensors='pt'
            )
            outputs = self.model(**inputs)
            scores = outputs.logits.squeeze(-1).numpy()

        return scores

    def rerank(self, query, documents, k=5):
        """文書をリランキング"""
        scores = self.score(query, documents)
        ranked_indices = np.argsort(scores)[::-1][:k]

        return [
            {'document': documents[i], 'score': scores[i]}
            for i in ranked_indices
        ]

# 使用例
documents = [
    "RAGは検索拡張生成の略称です。",
    "Pythonは人気のプログラミング言語です。",
    "RAGシステムは外部知識を活用してLLMの回答精度を向上させます。",
    "機械学習にはさまざまなアルゴリズムがあります。",
    "検索拡張生成はハルシネーションを軽減する効果があります。",
]

reranker = CrossEncoderReranker()
query = "RAGとは何ですか？"

results = reranker.rerank(query, documents, k=3)

print(f"Query: {query}\n")
print("Reranked results:")
for i, r in enumerate(results, 1):
    print(f"  {i}. {r['document']} (score: {r['score']:.4f})")

import numpy as np
from transformers import AutoTokenizer, AutoModel
import torch

class SimpleColBERT:
    def __init__(self, model_name='bert-base-uncased'):
        """ColBERTの簡易実装"""
        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
        self.model = AutoModel.from_pretrained(model_name)
        self.model.eval()

    def encode(self, text):
        """テキストをトークンレベルの埋め込みに変換"""
        with torch.no_grad():
            inputs = self.tokenizer(
                text,
                padding=True,
                truncation=True,
                max_length=512,
                return_tensors='pt'
            )
            outputs = self.model(**inputs)
            # 全トークンの隠れ状態を取得
            embeddings = outputs.last_hidden_state.squeeze(0).numpy()

        return embeddings

    def maxsim_score(self, query_emb, doc_emb):
        """MaxSimスコアを計算"""
        # (m, d) x (n, d).T = (m, n)
        similarity_matrix = np.dot(query_emb, doc_emb.T)
        # 各クエリトークンの最大類似度を合計
        max_similarities = np.max(similarity_matrix, axis=1)
        return np.sum(max_similarities)

    def rerank(self, query, documents, k=5):
        """文書をリランキング"""
        query_emb = self.encode(query)

        scores = []
        for doc in documents:
            doc_emb = self.encode(doc)
            score = self.maxsim_score(query_emb, doc_emb)
            scores.append(score)

        scores = np.array(scores)
        ranked_indices = np.argsort(scores)[::-1][:k]

        return [
            {'document': documents[i], 'score': scores[i]}
            for i in ranked_indices
        ]

# 使用例
colbert = SimpleColBERT()

documents = [
    "Machine learning is a subset of artificial intelligence.",
    "Python is widely used for data science.",
    "Deep learning uses neural networks with many layers.",
    "Natural language processing handles text data.",
]

query = "What is machine learning?"
results = colbert.rerank(query, documents, k=3)

print(f"Query: {query}\n")
print("ColBERT reranked results:")
for i, r in enumerate(results, 1):
    print(f"  {i}. {r['document']} (score: {r['score']:.2f})")

def llm_pointwise_rerank(query, documents, llm_client):
    """LLMによるPointwiseリランキング（疑似コード）"""
    scores = []

    for doc in documents:
        prompt = f"""以下の文書が質問に対してどの程度関連しているか、0から10のスコアで評価してください。

質問: {query}
文書: {doc}

スコア（0-10の整数のみ）:"""

        response = llm_client.generate(prompt)
        score = int(response.strip())
        scores.append(score)

    # スコア順にソート
    ranked = sorted(zip(documents, scores), key=lambda x: -x[1])
    return ranked

def llm_pairwise_compare(query, doc_a, doc_b, llm_client):
    """2文書を比較"""
    prompt = f"""質問に対して、どちらの文書がより関連していますか？

質問: {query}

文書A: {doc_a}
文書B: {doc_b}

「A」または「B」で回答:"""

    response = llm_client.generate(prompt)
    return 'A' if 'A' in response else 'B'

def llm_pairwise_rerank(query, documents, llm_client):
    """LLMによるPairwiseリランキング（バブルソート）"""
    docs = documents.copy()

    for i in range(len(docs)):
        for j in range(len(docs) - i - 1):
            winner = llm_pairwise_compare(query, docs[j], docs[j+1], llm_client)
            if winner == 'B':
                docs[j], docs[j+1] = docs[j+1], docs[j]

    return docs

import numpy as np

def dcg_at_k(relevances, k):
    """DCG@kを計算"""
    relevances = np.array(relevances)[:k]
    gains = 2 ** relevances - 1
    discounts = np.log2(np.arange(len(relevances)) + 2)
    return np.sum(gains / discounts)

def ndcg_at_k(relevances, k):
    """NDCG@kを計算"""
    dcg = dcg_at_k(relevances, k)
    ideal_relevances = sorted(relevances, reverse=True)
    idcg = dcg_at_k(ideal_relevances, k)
    return dcg / idcg if idcg > 0 else 0

def mrr(ranks):
    """MRRを計算"""
    return np.mean([1.0 / r for r in ranks if r > 0])

# 使用例
# 関連度スコア（実際の順序）
relevances = [3, 2, 0, 1, 0]  # 1位が関連度3, 2位が関連度2, ...

print(f"NDCG@3: {ndcg_at_k(relevances, 3):.4f}")
print(f"NDCG@5: {ndcg_at_k(relevances, 5):.4f}")

# 正解文書の順位
ranks = [1, 3, 2, 5]  # 4クエリで正解が1位, 3位, 2位, 5位
print(f"MRR: {mrr(ranks):.4f}")