Checklist OAI 2025 Stage II · Task 3

Source Extraction

Polish title: Ekstrakcja Źródeł

Produce 768-dimensional embeddings of scientific claims and abstracts with a supplied SGPT (GPT-2-based) model so that cosine-similarity retrieval finds each claim's source on SciFact.

NLP
Dense retrieval (embedding design)
Polish original · English translation

Download the English notebook Open the task notebook (Polish)

The task

Retrieval-augmented systems first search a document collection and only then generate an answer, which makes answers verifiable if the right sources are found. The contestant builds a module for a scientific fact-checking tool that quickly finds publications supporting or refuting a given claim.

The contestant implements Embedder.encode_queries and Embedder.encode_corpus, which map queries and documents (title and abstract) to real vectors of dimension 768. The provided evaluation code ranks all documents by cosine similarity to each query and takes the k = 10 nearest neighbours. A GPT-2-based model fine-tuned for embeddings (Muennighoff/SGPT-125M-weightedmean-msmarco-specb-bitfit) may be used.

Evaluation uses the SciFact benchmark: claims based on real publications and a corpus of life-science and medical abstracts, with at least one supporting or refuting publication per claim. The test queries do not include the matching document identifiers.

Abridged and translated by SOTA from the official Polish materials. The official statement has the exact rules, and it wins wherever this summary differs.

In English

This task was published in Polish. SOTA translated it into English on 16 September 2026. Only the words changed in the notebooks: markdown, code comments, messages and printed output. The code, file names and paths are the original's, so a translated notebook runs with the original data.

Task notebookEnglish · .ipynb · 24 KB Polish original of Task notebook

Read the task notebook in English 1412 words and 8 code cells

Source Extraction

Image generated with ChatGPT.

Introduction

Language models are prone to telling untruths or half-truths, and to making up facts without citing sources. Increasingly, systems are now used that, instead of answering questions directly, first search a database, e.g. a collection of documents, and only then generate an answer on the basis of the best-matching documents. Such an answer is more likely to be grounded in reality and can be verified by a human -- provided that the right sources were found correctly.

Of course, there may be a very large number of sources, so the search methods must be efficient -- processing everything "at once" directly with a language model is out of the question! In this task you will focus on finding the best sources for a given sentence, using the method of embeddings (vector embeddings).

Imagine that you are an AI engineer at a company developing a tool for verifying scientific facts. Your task is to build a module that can quickly and effectively find reliable scientific publications that confirm or refute specific claims. Thanks to your solution, scientists, journalists and decision-makers will be able to verify information on a solid scientific basis, which is especially important in an age of disinformation.

Task

Your task is to develop a system that generates high-quality vector representations (embeddings) for both queries and source documents, enabling the right sources to be matched precisely to the queries.

Given the queries (the set of queries; the queries for which we are looking for sources) and the corpus (the base of documents/sources; the set of documents under consideration), you must implement functions that assign to the queries and to the sources vectors of real numbers of dimension $768$ . These vectors will be used to find sources for each query by the evaluation function we provide, which, for a given query, selects the $k=10$ nearest neighbours ( $k$ -Nearest Neighbours) from the set of documents.

In your solution you may use the provided model based on the GPT2 architecture, which has been specially fine-tuned to help in obtaining good-quality embeddings.

While working on your solution, you will be able to test its effectiveness on the validation set, which will let you assess the quality of the generated embeddings in the context of the task of retrieving the right source documents.

Data

The data available to you in this task are:

A set of queries (queries) for which the matching sources must be found
A corpus of documents (corpus) containing scientific publications that may be sources for the queries
Information on which documents match the queries in the validation set

Your solution will be evaluated on the SciFact benchmark. It is used to evaluate fact retrieval and verification systems in a scientific context. It consists of a set of claims (queries) based on real scientific publications, and the document base (corpus) consists of publications in the natural and medical sciences. For each claim there is at least one publication that supports or refutes it. We provide the code for loading the data, so the data are described here for information only.

The file corpus.jsonl contains unique identifiers, titles and abstracts of scientific papers

Example of a single document:

{
    "text_id": 13734012,
    "title": "Prevalent abnormal prion protein in human appendixes after bovine spongiform encephalopathy epizootic: large scale survey",
    "text": "OBJECTIVES To carry out a further survey (...) CONCLUSIONS This study corroborates previous studies and suggests a high prevalence of infection with abnormal PrP, indicating vCJD carrier status in the population compared with the 177 vCJD cases to date. These findings have important implications for the management of blood and blood products and for the handling of surgical instruments."
}

The file queries_val.jsonl contains the text of the claims and the identifier of the matching source text. The test set, on which your solution will ultimately be scored, will not contain the identifiers of the matching source texts.

Example of a single query:

{
    "query": "1 in 5 million in UK have abnormal PrP positivity.",
    "matching_text_id": 13734012
}

Scoring Criterion

The methods (functions) Embedder.encode_queries and Embedder.encode_corpus that you implement will be used to convert, respectively, the queries $q \in Q$ and the documents $d \in C$ into vectors. In what follows, we will use $q$ and $d$ interchangeably to refer both to the texts and to their embeddings.

Suppose that the query $q\in Q$ corresponds to the gold document $d\in C$ . The evaluation code sorts all documents by their distance from $q$ , obtaining the documents $K_1, K_2, ..., K_n$ , such that $K_1$ is the closest. We then denote by $I$ the index of the gold document $d$ in this sequence. This means that $I - 1$ is the number of documents whose distance from $q$ is smaller than the distance of $q$ from $d$ .

The distance between vectors is computed using cosine similarity, which for vectors $v, w \in \mathbb{R}^n$ is defined as $\frac{v^Tw}{||v|| \cdot ||w||}$ , where $||v||$ is the length of the vector $v$ .

The score for a query $q$ is defined as

\text{nDCG@10}(q) = \begin{cases} \frac{1}{\log_2(I + 1)} & \text{if $I \leq 10$} \\ 0 & \text{otherwise.} \end{cases}

That is, the closer to the query the gold document has been placed relative to the other documents, the higher the score -- if 10 "wrong" documents are closer to the query, the score for this example is 0.

Ultimately, the scoring of your solution will be based on the nDCG@10 metric, computed as the mean value of this metric over all queries $(q \in Q )$ .

If the nDCG@10 score is lower than 0.2, you will receive 0 points.
If the score exceeds 0.5, you will receive the maximum number of points, i.e. 100.

For values between these thresholds, points will be awarded proportionally.

Constraints

Your solution will be tested on the Contest Platform without internet access and in an environment with a GPU.
The evaluation of your final solution on the Contest Platform must not take longer than 10 minutes with a GPU.
The embedding of each query and of each text must have dimension 768
List of permitted libraries: torch, pandas, numpy, nltk, transformers.

Submission Files

Submit only this notebook, completed with your solution (see the Embedder class).

Hints

The GPT2 model is a decoder-type language model. Decoder-type models work as follows: for a given sequence of tokens (e.g. a prefix of the sentence being processed) $t_1, t_2, \dots, t_n$ they compute a hidden vector $h_{n+1} \in \mathbb{R}^d$ , and then transform it with one of their weight matrices into $p_{n+1} \in \mathbb{R}^m$ -- a probability distribution over the tokens in the vocabulary.
Compared with the available execution time, there are many documents.

Evaluation

During grading, the FINAL_EVALUATION_MODE flag will be set to True.

For this task you can score between 0 and 100 points. The number of points you score will be computed on the (secret) test set on the Contest Platform using the formula given above, rounded to an integer. If your solution does not meet the above criteria or does not run correctly, you will receive 0 points for the task.

Starter Code

In this section we initialise the environment by importing the required libraries and functions. The prepared code for the tokenizer, data loading and evaluation will make it easier for you to work with the data and will allow you to solve the task.

######################### DO NOT CHANGE THIS CELL ##########################

FINAL_EVALUATION_MODE = False  # While grading your solution, we will change this value to True

######################### DO NOT CHANGE THIS CELL ##########################

import json
import os
from math import log2

import torch
from tqdm import tqdm
from transformers import AutoModel, AutoTokenizer


device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

class Tokenizer:
    def __init__(self, tokenizer_path, length=150):
        self.tokenizer = AutoTokenizer.from_pretrained(tokenizer_path)
        self.tokenizer.pad_token = self.tokenizer.eos_token
        self.tokenizer.padding_side = "right"
        self.length = length

    def __call__(self, batch_text):
        batch_tensor = self.tokenizer(
            batch_text,
            max_length=self.length,
            truncation=True,
            padding=True,
            return_tensors="pt"
        )
        return batch_tensor.to(device)

Loading the Data

In this part of the task we will load the training data.

######################### DO NOT CHANGE THIS CELL ##########################

def load_corpus(file):
    corpus = {}
    with open(file, encoding="utf8") as f_in:
        for line in f_in:
            line = json.loads(line)
            corpus[line.get("text_id")] = {
                "text": line.get("text"),
                "title": line.get("title"),
            }
    return corpus

def load_queries(file):
    queries = {}
    matching_texts = {}
    with open(file, encoding="utf8") as f_in:
        for query_num, line in enumerate(f_in):
            line = json.loads(line)

            queries[query_num] = line.get("query")
            matching_texts[query_num] = line.get("matching_text_id")
    return queries, matching_texts

corpus = load_corpus("corpus.jsonl")
queries, matching_texts = load_queries("queries_val.jsonl")

print(f"Loaded {len(corpus)} texts and {len(queries)} queries.")

Scoring Criterion Code

Code similar to the code below will be used to evaluate the solution on the test set.

######################### DO NOT CHANGE THIS CELL ##########################

def evaluate_retrieval_ndcg(
    golden_matches: dict[int, int],
    results: dict[int, dict[int, float]],
) -> float:
    """
    Computes the value of the nDCG metric for the given search results.

    The function computes the score of your solution based on the results for the $top\\_k$ best documents according to your embedder.

    :param golden_matches: Dictionary with the gold assignments, where the key is the query id and the value is the id of the correct document.
    :param results: Dictionary with the search results, where the key is the query id and the value is a dictionary of document ids and their similarities to the given query.
    :return: The value of the nDCG metric."""

    for query_id, v in results.items():
        results[query_id] = {k: v for k, v in sorted(v.items(), key=lambda item: -item[1])}

    ndcg_sum = 0
    for query_id, v in results.items():
        golden_document = golden_matches[query_id]
        for i, document_id in enumerate(v.keys()):
            if golden_document == document_id:
                ndcg_sum += 1 / log2(i + 2)

    ndcg = round(ndcg_sum / len(results), 5)
    return ndcg


def compute_score(ndcg: float) -> float:
    """
    Computes the points score from the value of the nDCG metric.
    """
    lower_bound = 0.2
    upper_bound = 0.5

    if ndcg <= lower_bound:
        return 0
    elif lower_bound < ndcg < upper_bound:
        return int(round(100 * (ndcg - lower_bound) / (upper_bound - lower_bound)))
    else:
        return 100

Search

Below is the code used to select, for a given query, the $top\_k$ best documents from the corpus.

######################### DO NOT CHANGE THIS CELL ##########################

def cos_sim(a: torch.Tensor, b: torch.Tensor):
    """
    Computes the cosine similarity cos_sim(a[i], b[j]) for all i and j.
    :return: Matrix with res[i][j]  = cos_sim(a[i], b[j])
    """
    a_norm = torch.nn.functional.normalize(a, p=2, dim=1)
    b_norm = torch.nn.functional.normalize(b, p=2, dim=1)
    return torch.mm(a_norm, b_norm.transpose(0, 1))

def search_topk_texts(
    embedder,
    corpus: dict[str, dict[str, str]],
    queries: dict[str, str],
    top_k: int = 10,
) -> dict[str, dict[str, float]]:
    results = {}

    # Create embeddings for all queries using model.encode_queries()
    # Runs semantic search against the corpus embeddings
    # Returns a ranked list with the corpus ids
    query_ids = list(queries.keys())
    results = {qid: {} for qid in query_ids}
    queries = [queries[qid] for qid in queries]
    query_embeddings = embedder.encode_queries(queries)

    corpus_ids = sorted(
        corpus,
        key=lambda k: len(corpus[k].get("title", "") + corpus[k].get("text", "")),
        reverse=True,
    )
    corpus = [corpus[cid] for cid in corpus_ids]

    # Encode chunk of corpus
    corpus_embeddings = embedder.encode_corpus(corpus)

    # Compute similarites using cosine-similarity
    cos_scores = cos_sim(query_embeddings, corpus_embeddings)
    cos_scores[torch.isnan(cos_scores)] = -1

    # Get top-k values
    cos_scores_top_k_values, cos_scores_top_k_idx = torch.topk(
        cos_scores,
        min(top_k + 1, len(cos_scores[1])),
        dim=1,
        largest=True,
        sorted=False,
    )
    cos_scores_top_k_values = cos_scores_top_k_values.cpu().tolist()
    cos_scores_top_k_idx = cos_scores_top_k_idx.cpu().tolist()

    for query_itr in range(len(query_embeddings)):
        query_id = query_ids[query_itr]
        for score, corpus_id in zip(cos_scores_top_k_values[query_itr], cos_scores_top_k_idx[query_itr]):
            results[query_id][corpus_ids[corpus_id]] = score

    return results

Your Solution

Put your solution in this section. Make changes only here!

class Embedder:
    # Do not change the constructor's signature
    def __init__(self):
        # TODO: you may change this method,
        # but do not change its signature! (i.e. do not change the arguments)
        self.model = AutoModel.from_pretrained("Muennighoff/SGPT-125M-weightedmean-msmarco-specb-bitfit")
        self.tokenizer = Tokenizer("Muennighoff/SGPT-125M-weightedmean-msmarco-specb-bitfit")

    def encode_queries(self, queries: list[str]):
        """
        Function that encodes the queries.
        :param queries: List of queries to encode
        :return: Query embeddings - a tensor of shape (n, 768), where n = len(queries) is the number of queries.
        """

        # TODO: implement this method - encode the queries
        # Do not change this method's signature! (i.e. do not change the arguments)
        # Remember that you can use the HuggingFace gpt-2 model...
        # You can use the Tokenizer implemented in the upper part of the notebook
        # Hint: The evaluation will be faster if the returned tensor is on the GPU.
        ...
        return torch.ones(len(queries), 768).to(device)

    def encode_corpus(self, texts: list[dict]):
        """
        Function that encodes the source texts.
        :param texts: List of texts to encode. Each text is represented as a dictionary:
            {
                "title": "..."
                "text": "...",
            }
        :return: Text embeddings - a tensor of shape (m, 768), where m = len(texts) is the number of texts
        """

        # TODO: implement this method - encode the source texts
        # Do not change this method's signature! (i.e. do not change the arguments)
        ...
        return torch.ones(len(texts), 768).to(device)

Evaluation

Running the cell below lets you check how many points your solution would score on the validation data. Before submitting, make sure that the whole notebook runs from start to finish without errors and without any user intervention after selecting "Run All".

######################### DO NOT CHANGE THIS CELL ##########################

if not FINAL_EVALUATION_MODE:
    embedder = Embedder()

    with torch.no_grad():
        results = search_topk_texts(embedder, corpus, queries, top_k=10)

    # Compute nDCG
    ndcg = evaluate_retrieval_ndcg(matching_texts, results)

    # Compute the final score from nDCG
    points = compute_score(ndcg)

    print(f"\nNumber of queries: {len(queries)}")
    print(f"Number of texts: {len(corpus)}")
    print(f"nDCG: {ndcg:.3f}")
    print(f"Points score: {points}")

During grading, the model will be saved as your_model.pkl and scored on the test set.

######################### DO NOT CHANGE THIS CELL ##########################

if FINAL_EVALUATION_MODE:
    import cloudpickle

    OUTPUT_PATH = "file_output"
    FUNCTION_FILENAME = "your_model.pkl"
    FUNCTION_OUTPUT_PATH = os.path.join(OUTPUT_PATH, FUNCTION_FILENAME)

    if not os.path.exists(OUTPUT_PATH):
        os.makedirs(OUTPUT_PATH)

    with open(FUNCTION_OUTPUT_PATH, "wb") as f:
        cloudpickle.dump(Embedder, f)

Translated by SOTA. The Polish original is the official version and wins wherever the two differ. The notebook reads corpus.jsonl and queries_val.jsonl from its own folder; the original does not link to the data files. If you organise this olympiad and would like the translation removed, email [email protected] and we will take it down.

At a glance

You get: corpus.jsonl (text_id, title, text) and queries_val.jsonl (query, matching_text_id) in the task folder.
You submit: This notebook with the Embedder class; saved to your_model.pkl during checking.
Scoring: nDCG@10 per query = 1/log2(I + 1) if the gold document's rank I ≤ 10, else 0; averaged over queries. 0 points if nDCG@10 < 0.2, 100 points if above 0.5, linear in between.
Rules: Tested without Internet access, with a GPU; evaluation must take at most 10 minutes with a GPU.
Every query and document embedding must have dimension 768.
Allowed libraries: torch, pandas, numpy, nltk, transformers.
Hint: the corpus is large relative to the available time.
Format: Stage II (regional, on site in Kraków, Poznań, Warsaw and Wrocław, identical tasks in all cities), 26–27 April 2025; Evaluated automatically on the Competition Platform (Platforma Konkursowa) on a hidden test set; points are rounded to an integer, and a notebook that fails the requirements or does not run scores 0.

Details

Olympiad: Polish Olympiad in Artificial Intelligence (Olimpiada Sztucznej Inteligencji)
Year: 2025, Kraków, Poznań, Warsaw and Wrocław, Poland
Round: Stage II · Task 3
Language: Polish; English translation by SOTA
License: Not stated by the source

Source Extraction

The task

In English

Source Extraction

Introduction

Task

Data

Scoring Criterion

Constraints

Submission Files

Hints

Evaluation

Starter Code

Loading the Data

Scoring Criterion Code

Search

Your Solution

Evaluation

At a glance

Links

In English (SOTA)

Notebook

Data

Model

Solution

Results

Source files

Details