hotchpotch/mmarco-hard-negatives-reranker-score

target_languages = [
    "english",
    "chinese", 
    "french",
    "german",
    "indonesian",
    "italian",
    "portuguese",
    "russian",
    "spanish",
    "arabic",
    "dutch",
    "hindi",
    "japanese",
    "vietnamese"
]

# target languages => ["english", "chinese", "french", "german", "indonesian", "italian", "portuguese", "russian", "spanish", "arabic", "dutch", "hindi", "japanese", "vietnamese"]

lang = "spanish"
repo_id = "hotchpotch/mmarco-hard-negatives-reranker-score"
reranker = "bge-reranker-v2-m3"
subset = f"{lang}_{reranker}"
mapping = f"mappings/{lang}_joblib.pkl.gz"

from datasets import load_dataset
import joblib
from huggingface_hub import hf_hub_download

queries_ds = load_dataset(
    "unicamp-dl/mmarco", "queries-" + lang, split="train", trust_remote_code=True
)
collection_ds = load_dataset(
    "unicamp-dl/mmarco",
    "collection-" + lang,
    split="collection",
    trust_remote_code=True,
)
score_ds = load_dataset(repo_id, subset, split="train")


mapping_file = hf_hub_download(repo_type="dataset", repo_id=repo_id, filename=mapping)

index_mapping_dict = joblib.load(mapping_file)
query_id_dict = index_mapping_dict["query_id_dict"]
collection_id_dict = index_mapping_dict["collection_id_dict"]


def get_query_text(query_id) -> str:
    idx = query_id_dict[query_id]
    return queries_ds[idx]["text"]  # type: ignore


def get_collection_text(doc_id) -> str:
    idx = collection_id_dict[doc_id]
    return collection_ds[idx]["text"]  # type: ignore

for i in range(5):
    qid: int = score_ds[i]["qid"]
    pos: list[int] = score_ds[i]["pos"]
    pos_score: list[float] = score_ds[i]["pos.score"]
    neg: list[int] = score_ds[i]["neg"]
    neg_score: list[float] = score_ds[i]["neg.score"]
    query = get_query_text(qid)
    pos_docs = [get_collection_text(doc_id)[0:64] for doc_id in pos]
    neg_docs = [get_collection_text(doc_id)[0:64] for doc_id in neg]
    print(f"# Query: {query}")
    print("## Positive docs:")
    for doc, score in zip(pos_docs, pos_score):
        print(f"  {doc} ({score})")
    print("## Negative docs:")
    for doc, score in list(zip(neg_docs, neg_score))[0:5]:
        print(f"  {doc} ({score})")
    print("-------")

# Query: ¿Qué son las artes liberales?
## Positive docs:
  Artes liberales. 1. el curso académico de instrucción en una uni (0.99770385)
## Negative docs:
  Grandes Ligas. Puede elegir entre una variedad de especializacio (0.69760895)
  BA = Licenciatura en Artes BS = Licenciatura en Ciencias Creo qu (0.24364243)
  ¿Qué es una Licenciatura en Artes (B.A.)? Un programa de licenci (0.20641373)
  ¿Qué significa LCSW? / Human and Social ... / Liberal Arts y ... (0.0140636265)
  definición de artes liberales Las áreas de aprendizaje que culti (0.9963924)
-------
# Query: ¿Cuál es el mecanismo de acción de los fármacos fibrinolíticos o trombolíticos?
## Positive docs:
  Hematología clínica de BailliÃÆ'¨re. 6 Mecanismo de acción d (0.966347)
## Negative docs:
  Definición y ejemplos de mecanismos de acción. Más en Trastorno  (0.3598139)
  ¿Qué es losartán y cómo funciona (mecanismo de acción)? ¿Qué mar (0.0031480708)
  ActivaseÃ⠀ šÃ,® Una propietaria trombolítico, que puede à ¢  (0.83237296)
  La terapia fibrinolítica, también llamada a veces "terapia tromb (0.92590266)
  Diazepam Valium Mecanismo de acción Valium Mecanismo de acción E (0.040162582)
-------
# Query: ¿Qué es el recuento normal de plataformas?
## Positive docs:
  78 seguidores. R. Las plaquetas son glóbulos diminutos que ayuda (0.10105592)
## Negative docs:
  ¿Qué es la trombocitopenia (recuento bajo de plaquetas)? Las pla (0.047337715)
  Calificación Más reciente Más antiguo. Mejor respuesta: Nancy: e (0.03179867)
  1 Tarifas de solicitud ࢠ€Â⠀ œ $ 80 por una plataform (0.001044386)
  Conteo sanguíneo de MCV. Mi recuento sanguíneo de MCV está en 98 (0.011115014)
  ¿Cuáles son los niveles normales de hemograma para una mujer adu (0.015247591)
-------
# Query: promedio de costo en dólares explicado
## Positive docs:
  El promedio del costo en dólares es simplemente un método para c (0.96771675)
## Negative docs:
  El promedio del costo en dólares es una técnica simple que le pe (0.9859364)
  Anteriormente en Free from Broke, Glen ha abordado el tema del c (0.89274967)
  (TMFMathGuy). 19 de noviembre de 2014 a las 2:15 p.m. El promedi (0.98421544)
  Comprar acciones por valor de $ 2,000 en el primer mes, a $ 14.2 (0.36624968)
  DEFINICIÓN de 'Valor Promedio'. Una estrategia de inversión que  (0.3941427)
-------
# Query: alimentos que ayudan a combatir la gota
## Positive docs:
  Además de seguir una dieta bien balanceada para promover la salu (0.9874721)
## Negative docs:
  Alimentos para la gota: anacardos y gota. Los anacardos y la got (0.84606963)
  Alimentos que debe evitar si tiene gota La gota es un tipo de ar (0.016979992)
  El puerto causa gota. Estos son solo algunos tratamientos natura (0.30850264)
  20 alimentos ricos en purina que debe evitar: alterar su dieta p (0.4324828)
  Alimentos que combaten el cáncer por BÃÆ' © liveau & Gingras (0.010408314)
-------