generate_retrival_metrisc_using_beir-STEP-1.py

import os
import pathlib
import logging
from beir import LoggingHandler
from beir.datasets.data_loader import GenericDataLoader
from beir.retrieval.evaluation import EvaluateRetrieval
from sentence_transformers import SentenceTransformer
import chromadb
from langchain_core.prompts import ChatPromptTemplate
from pydantic import BaseModel, Field
from llm_client import together_client
from modul_decyzujny.decompose_query import check_if_query_is_complex, decompose_query
from modul_decyzujny.first_router import query_router, ToolChoice
from modul_decyzujny.knowledge_summarizer import (
    summarize_knowledge_with_most_occurring_words,
)
from query_analizer.hyde_generator import HyDEGenerator

from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch

import re
import json

from tools.vector_store_search_fallback import vector_store_search

# use llm as a judge to check if document is relevant to question
# mozna uzywac llm ALBO rerankera

# ============== parameters to adjust ===================
# specify dataset HERE
# fiqa - ma normalne pytania
# trec-covid tez git

# dataset = "trec-covid-v2"
dataset = "fiqa"
use_ll_aaj = False
use_reranker = False
use_hyde = False
split_documents = False
route_questions = False
use_adaptive_k = False
# moze byc tu odpalic mistrall small 3.1
judge_model = "mistralai/Mixtral-8x7B-Instruct-v0.1"
# cross_encoder_model = "cross-encoder/ms-marco-MiniLM-L-6-v2"
cross_encoder_model = "cross-encoder/quora-distilroberta-base"
# ========================================================


assert not (
    use_ll_aaj and use_reranker
), "Ustaw tylko jeden tryb oceny: LLM albo reranker!"


if use_reranker:
    cross_encoder_name = cross_encoder_model
    ce_tokenizer = AutoTokenizer.from_pretrained(cross_encoder_name)
    ce_model = AutoModelForSequenceClassification.from_pretrained(cross_encoder_name)
    ce_model.eval()


def grade_document_with_mistral(query: str, document: str) -> str:
    grading_prompt = f"""
    You are a grader assessing relevance of a document to a question.

    You are given:
    1. A question
    2. Part of document, that might be helpfully to answer this question. 

    If this document might be helpful in answering question that you will see in a second, or it has  semantic meaning 
    related to the question, it should be considered as relevant otherwise it is not relevant. 
    Your task is to provide answer that will inform if document is highly relevant to the question or not.  

    Now, you will question and part of this document.

    ---
    Question:
    {query}

    Document:
    {document}

    Please respond ONLY with a JSON object like this, and make aure your judge is relevant to previous instructions:
    {{"response": "yes"}} or {{"response": "no"}}

    """

    response = together_client.chat.completions.create(
        model=judge_model,
        messages=[{"role": "user", "content": grading_prompt}],
        temperature=0,
        max_tokens=50,
    )

    response_text = response.choices[0].message.content.strip()

    try:
        match = re.search(r"\{.*?\}", response_text)
        if not match:
            raise ValueError("No JSON object found in response.")
        parsed = json.loads(match.group(0))
        print("✅ Successfully parsed json")

        return parsed.get("response", "").strip().lower()

    except Exception as e:
        print(f"❌ JSON parsing failed: {e}\nResponse was: {response_text}")
        return "no"  # default fallback


class GradeDocuments(BaseModel):
    binary_score: str = Field(
        description="Documents are relevant to the question, 'yes' or 'no'"
    )


preamble = """You are a grader assessing relevance of a retrieved document to a user question.
If the document contains keyword(s) or semantic meaning related to the user question, grade it as relevant.
Give a binary score 'yes' or 'no' score to indicate whether the document is relevant to the question."""

grade_prompt = ChatPromptTemplate.from_messages(
    [("human", "Retrieved document: \n\n{document}\n\nUser question: {question}")]
)

# === Logowanie ===
logging.basicConfig(
    format="%(asctime)s - %(message)s",
    datefmt="%Y-%m-%d %H:%M:%S",
    level=logging.INFO,
    handlers=[LoggingHandler()],
)

# === Parametry ===

embedding_model_name = "all-MiniLM-L6-v2"
collection_name = f"beir-{dataset}-corpus"
top_k = 32

# === 1. Załaduj dane ===
data_path = os.path.join(pathlib.Path(__file__).parent.absolute(), "datasets", dataset)
corpus, queries, qrels = GenericDataLoader(data_folder=data_path).load(split="test")


if route_questions:
    docs_summary = summarize_knowledge_with_most_occurring_words(corpus, top_k=20)

# === 2. Model ===
model = SentenceTransformer(embedding_model_name)


# inizjaliazcja modułu hyde
if use_hyde:
    hyde_gen = HyDEGenerator(
        # generation_model="mistralai/Mistral-7B-Instruct-v0.1",
        generation_model=judge_model,
        embedding_model=embedding_model_name,
    )

# === 3. Chroma ===
client = chromadb.PersistentClient(path="chroma_store")
collection = client.get_collection(name=collection_name)

# results to jest slownik, gdzie jako klucz mamy id pytania, a jako wartosc mamy kolejny wlonik w ktorym jes twiele kluczy i warotsci
# te klucze to iq pobranych dokumentow, a warotsci to ich score
results = {}

logging.info("Wyszukiwanie przez Chroma + budowanie results...")

number_of_questions_to_process = 0
number_of_irrelevant_answers = 0


for qid, query_text in queries.items():
    if number_of_questions_to_process == 20:
        break

    all_hits = []

    print(f"\nCurrent query: {query_text}")

    if use_adaptive_k:
        query_hits, k = vector_store_search(
            query_text, embedding_model=model, collection=collection, top_k=32
        )

    else:

        query_emb = model.encode(query_text, convert_to_numpy=True).tolist()
        query_hits = collection.query(query_embeddings=[query_emb], n_results=top_k)

    for doc_id, dist in zip(query_hits["ids"][0], query_hits["distances"][0]):
        all_hits.append((doc_id, dist))

    if use_hyde:
        try:
            hyde_emb = hyde_gen.embed_hyde(query_text)
        except Exception as e:
            logging.warning(f"HyDE generation failed for query {qid}: {e}")
            hyde_emb = []
        hits_hyde = (
            collection.query(query_embeddings=[hyde_emb], n_results=top_k)
            if hyde_emb
            else {"ids": [[]], "distances": [[]]}
        )

    if split_documents:
        is_query_complex = check_if_query_is_complex(query=query_text)

        if is_query_complex:
            subquestions = decompose_query(query=query_text)
            for sq in subquestions:
                print(f" - {sq}")
        else:
            subquestions = [query_text]

        for sub_query in subquestions:

            # dodac wazony scoring
            # top_k_per_subq = round(top_k / len(subquestions) + 2)
            top_k_per_subq = 5

            if route_questions:
                response_from_router = query_router(
                    query_text=sub_query, docs_summary=docs_summary
                )
                if response_from_router == ToolChoice.VECTORSTORE:

                    subquery_embeddings = model.encode(
                        sub_query, convert_to_numpy=True
                    ).tolist()
                    subquery_hits = collection.query(
                        query_embeddings=[subquery_embeddings], n_results=top_k_per_subq
                    )

                    for doc_id, dist in zip(
                        subquery_hits["ids"][0], subquery_hits["distances"][0]
                    ):
                        all_hits.append((doc_id, dist))

            else:
                # metody na poprawienie metryk podczas rozbicia
                # zmiana top_k
                # pobranie tylko np dwoch pierwszych podpytan
                # uzycie scoringu wazonego
                # dodac wazony score w zaleznosci od źrodła, bardziej licza sie dokumenty z orginalnego pytania
                # pozniej mozna brac jescze topkumenyt o danym top scorze
                subquery_embeddings = model.encode(
                    sub_query, convert_to_numpy=True
                ).tolist()
                subquery_hits = collection.query(
                    query_embeddings=[subquery_embeddings], n_results=top_k_per_subq
                )

                for doc_id, dist in zip(
                    subquery_hits["ids"][0], subquery_hits["distances"][0]
                ):
                    all_hits.append((doc_id, dist))

    doc_scores = {}

    all_hits = sorted(all_hits, key=lambda x: x[1])[:top_k]

    for doc_id, dist in all_hits:
        score = 1 / (1 + dist)
        if doc_id in doc_scores:
            doc_scores[doc_id] = max(doc_scores[doc_id], score)
        else:
            doc_scores[doc_id] = score

    doc_ids = list(doc_scores.keys())
    scores = list(doc_scores.values())

    if use_ll_aaj:
        results[qid] = {}
        for doc_id, score in zip(doc_ids, scores):
            doc_txt = corpus[doc_id]

            try:
                doc_text = doc_txt.get("text", "")
                grade = grade_document_with_mistral(query_text, doc_text)
                if grade == "yes":
                    results[qid][doc_id] = score
                if grade == "no":
                    number_of_irrelevant_answers += 1
            except Exception as e:
                logging.warning(f"LLM grading failed for doc {doc_id}: {e}")
    elif use_reranker:
        pairs = [(query_text, corpus[doc_id].get("text", "")) for doc_id in doc_ids]

        # here we tokenize and predict
        with torch.no_grad():
            inputs = ce_tokenizer.batch_encode_plus(
                pairs, padding=True, truncation=True, return_tensors="pt"
            )
            scores_cs = ce_model(**inputs).logits.squeeze().tolist()

        if isinstance(scores_cs, float):
            scores_cs = [scores_cs]

        results[qid] = dict(
            sorted(
                ((doc_id, float(score)) for doc_id, score in zip(doc_ids, scores_cs)),
                key=lambda x: x[1],
                reverse=True,
            )
        )
    else:
        results[qid] = dict(
            sorted(doc_scores.items(), key=lambda x: x[1], reverse=True)
        )

    number_of_questions_to_process += 1

print(f"Ilosc odrzuconych odpowiedzi: {number_of_irrelevant_answers}")
logging.info(f"Gotowe — oceniam {len(results)} zapytań.")

# === 5. Ocena BEIR ===
evaluator = EvaluateRetrieval()

number_of_k_values_to_check = [4, 8, 16, 32]

ndcg, _map, recall, precision = evaluator.evaluate(
    qrels, results, number_of_k_values_to_check
)

print("\n📊 Wyniki oceny:")
for k in number_of_k_values_to_check:
    print(f"NDCG@{k}: {ndcg[f'NDCG@{k}']:.4f}")
    print(f"MAP@{k}: {_map[f'MAP@{k}']:.4f}")
    print(f"Recall@{k}: {recall[f'Recall@{k}']:.4f}")
    print(f"P@{k}: {precision[f'P@{k}']:.4f}")
    print()

# === 6. Zapisz wyniki do JSON ===
results_dir = os.path.join(pathlib.Path(__file__).parent.absolute(), "results")
os.makedirs(results_dir, exist_ok=True)

# Sklej wszystkie metryki do jednego słownika
full_metrics = {"NDCG": ndcg, "MAP": _map, "Recall": recall, "Precision": precision}

suffix_name = ""

if use_ll_aaj:
    suffix_name += "_judged_by_llm"
elif use_reranker:
    model_name_to_save_to_file = cross_encoder_model.replace("/", "-")
    suffix_name += f"_reranked_{model_name_to_save_to_file}"

if use_hyde:
    suffix_name += "_and_hyde"

if split_documents:
    suffix_name += "_questions_splitted"

if use_adaptive_k:
    suffix_name += "_adaptive_k_applied"

output_path = os.path.join(
    results_dir, f"fixed_{dataset}-{embedding_model_name}_first_20_q{suffix_name}.json"
)

with open(output_path, "w") as f:
    json.dump(full_metrics, f, indent=4)

print(f"✅ Wyniki zapisane do pliku: {output_path}")