instruct_model.py

from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    PreTrainedModel,
    PreTrainedTokenizer,
)
import torch
import torch.nn.functional as F
import json
import logging
from fuzzysearch import find_near_matches
from typing import Mapping
from utils import json_kwargs


INSTRUCT_MODEL = "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B"  # Don't use it - it forces the reasoning process, which we don't want here
INSTRUCT_MODEL = "Qwen/Qwen2.5-1.5B-Instruct"
INSTRUCT_MODEL = "BSC-LT/salamandra-7b-instruct"
INSTRUCT_MODEL = "utter-project/EuroLLM-9B-Instruct"
INSTRUCT_MODEL = "CohereForAI/aya-expanse-8b"
INSTRUCT_MODEL = "Qwen/Qwen2.5-7B-Instruct"
INSTRUCT_MODEL = "meta-llama/Llama-3.1-8B-Instruct"

instruct_model_tokenizer = AutoTokenizer.from_pretrained(
    INSTRUCT_MODEL, local_files_only=False, trust_remote_code=True
)
# Set the padding token to be the same as the eos token
instruct_model_tokenizer.pad_token = instruct_model_tokenizer.eos_token
instruct_model: PreTrainedModel = AutoModelForCausalLM.from_pretrained(
    INSTRUCT_MODEL,
    device_map="auto",
    trust_remote_code=True,
    local_files_only=False,
)
instruct_model.eval()


@torch.no_grad()
def tag_hallucination_spans_words(
    words_hint: str,
    question: str,
    answer: str,
    instruct_model: PreTrainedModel,
    instruct_model_tokenizer: PreTrainedTokenizer,
):
    """
    Tags the hallucination spans in the full text taking into account the word hint.
    """
    tokenized_text = instruct_model_tokenizer.apply_chat_template(
        conversation=[
            {
                "role": "system",
                "content": f'You are an annotator of spans in a text. The user will provide you with a hint about an idea that you should extract in the text. You should return the minimal span of the text that is a complete semantic unit, such as an entity or a clause, that is most likely to be related to the hint. Always produce a result, even if you need to make assumptions. Write the minimal span of text that you select exactly as it appeared in the text, do not add or remove anything or include any markers or special characters. Do not output any errors or information, exclusively output the result. Follow exactly this JSON schema and don\'t output anything else: {{"extracted_entity": "The minimal span of the text that is most likely to be related to the hint, as it appears in the text."}}.',
            },
            {
                "role": "user",
                "content": json.dumps(
                    {
                        "question": "¿Quién ganó la Copa de Clubes de Asia 1997-98?",
                        "text": "La Copa de Clubes de Asia 1997-98 fue ganada por el club de fútbol japonés, el Kashima Antlers. Derrotaron al club de fútbol tailandés, el BEC Tero Sasana, 1-0 en la final, que se celebró en el Estadio Nacional de la ciudad de Bangkok, Tailandia, el 31 de octubre de 1998.",
                        "hint": "japonés, el",
                    },
                    **json_kwargs,
                ),
            },
            {
                "role": "assistant",
                "content": json.dumps(
                    {"extracted_entity": "Kashima Antlers"}, **json_kwargs
                ),
            },
            {
                "role": "user",
                "content": json.dumps(
                    {
                        "question": "Est-ce que Chance the Rapper a joué au Stade de France?",
                        "text": "Oui, Chance the Rapper a joué au Stade de France lors de son concert le 23 juin 2017. Il a également joué au Stade de France lors de la cérémonie d'ouverture des Jeux olympiques de la jeunesse de Paris 2018 le 6 octobre 2018.",
                        "hint": "olympiques de",
                    },
                    **json_kwargs,
                ),
            },
            {
                "role": "assistant",
                "content": json.dumps(
                    {
                        "extracted_entity": "Jeux olympiques de la jeunesse de Paris 2018"
                    },
                    **json_kwargs,
                ),
            },
            {
                "role": "user",
                "content": json.dumps(
                    {
                        "question": "¿Cuál es un dato curioso sobre Beethoven?",
                        "text": "Beethoven es el único compositor de música clásica que no llegó a escribir ninguna obra mientras estaba vivo.",
                        "hint": "que no llegó",
                    },
                    **json_kwargs,
                ),
            },
            {
                "role": "assistant",
                "content": json.dumps(
                    {"extracted_entity": "no llegó"},
                    **json_kwargs,
                ),
            },
            {
                "role": "user",
                "content": json.dumps(
                    {
                        "question": "If I have 17 apples and eat 2, how many apples do I have left?",
                        "text": "If you eat 2, you will have 15 apples left.",
                        "hint": "have 15 apples",
                    },
                    **json_kwargs,
                ),
            },
            {
                "role": "assistant",
                "content": json.dumps(
                    {"extracted_entity": "15"},
                    **json_kwargs,
                ),
            },
            {
                "role": "user",
                "content": json.dumps(
                    {
                        "question": "Wer spielt Dexter in der Fernsehserie?",
                        "text": "Dexter wird von Michael Hall gespielt.",
                        "hint": "wird von Michael",
                    },
                    **json_kwargs,
                ),
            },
            {
                "role": "assistant",
                "content": json.dumps(
                    {"extracted_entity": "Michael Hall"},
                    **json_kwargs,
                ),
            },
            {
                "role": "user",
                "content": json.dumps(
                    {
                        "question": question,
                        "text": answer,
                        "hint": words_hint,
                    },
                    **json_kwargs,
                ),
            },
            {
                "role": "assistant",
                "content": '{"extracted_entity":',
            },
        ],
        add_generation_prompt=False,
        continue_final_message=True,
        return_dict=True,
        return_tensors="pt",
    )

    # Run the generation
    instruct_model_outputs = instruct_model.generate(
        tokenized_text["input_ids"].to(instruct_model.device),
        attention_mask=tokenized_text["attention_mask"].to(instruct_model.device),
        max_new_tokens=256,
        num_beams=3,
        early_stopping=True,
        # length_penalty=-1.0,
        do_sample=True,
        # top_p=0.95,
        # temperature=1.0,
        num_return_sequences=1,
    )

    logging.debug(
        f"Prompt: {instruct_model_tokenizer.decode(tokenized_text['input_ids'][0])}"
    )

    # Remove the part of the tokenized text that is the prompt... except for the part where we start the JSON, more or less 8 tokens from the end
    instruct_model_outputs = instruct_model_outputs[
        ..., tokenized_text["input_ids"].size(-1) - 7 :
    ].squeeze()

    logging.debug(f"Instruct model outputs: {instruct_model_outputs}")

    # Decode the tokenized text
    decoded_text = instruct_model_tokenizer.decode(
        instruct_model_outputs, skip_special_tokens=True
    )

    logging.debug(f"Decoded text (hallucination spans): {decoded_text}")

    # Remove everything before the opening bracket
    decoded_text = decoded_text[decoded_text.index("{") :]

    # Remove everything after the last closing bracket
    decoded_text = decoded_text[: decoded_text.rindex("}") + 1]

    logging.debug(
        f"Decoded text (hallucination spans, after adjustment): {decoded_text}"
    )

    lower_full_text = answer.lower()

    try:
        decoded_json = json.loads(decoded_text)

        extracted_entity_text = decoded_json["extracted_entity"]

        lower_extracted_entity_text = extracted_entity_text.lower()

        assert extracted_entity_text, f"Extracted span is empty: {decoded_json}"

        # Offsets is a tuple (start_char, end_char)
        # Use fuzzysearch to find the extracted entity in the full text, with a maximum Levenshtein distance of 20% of the length of the extracted entity
        extracted_entity_matches = find_near_matches(
            lower_extracted_entity_text,
            lower_full_text,
            max_l_dist=(len(lower_extracted_entity_text) // 5),
        )
        if not extracted_entity_matches:
            raise ValueError(
                f"Could not find the extracted entity in the full text: {extracted_entity_text}"
            )

        # Get the first match
        extracted_entity_match = extracted_entity_matches[0]
        extracted_entity_text = answer[
            extracted_entity_match.start : extracted_entity_match.end
        ]

        # Get the offsets
        offsets = (extracted_entity_match.start, extracted_entity_match.end)

    except (KeyError, json.JSONDecodeError, ValueError) as e:
        logging.error(f'Error decoding JSON: "{decoded_text}": {e}')
        extracted_entity_text = words_hint
        lower_extracted_entity_text = extracted_entity_text.lower()

        offsets = (
            lower_full_text.index(lower_extracted_entity_text),
            lower_full_text.index(lower_extracted_entity_text)
            + len(lower_extracted_entity_text),
        )

    # We can't have a finishing offset == len(full_text) because that would be out of bounds. If we do, we'll subtract 1. Don't do it if we only have one character - this is a kludge but don't have time to fix it properly now
    if offsets[1] == len(lower_full_text) and len(lower_full_text) > 1:
        offsets = (offsets[0], offsets[1] - 1)
        logging.warning(
            f"Adjusted offsets because the end was at the end of the text: {offsets}"
        )

    # The end offset must be at least the start offset + 1
    if offsets[1] <= offsets[0]:
        logging.warning(
            f"Adjusted offsets because the end was before the start: {offsets}"
        )
        offsets = (offsets[0], offsets[0] + 1)

    logging.info(
        f'Extracted entity "{words_hint}": -> "{extracted_entity_text}" (from "{answer}")'
    )

    return extracted_entity_text, offsets


@torch.no_grad()
def assign_priorities_to_models(
    question: str,
    answer: str,
    weight_temperature,
    instruct_model: PreTrainedModel,
    instruct_model_tokenizer: PreTrainedTokenizer,
    languages_to_models: Mapping[str, str],
):
    """
    This function asks an English LLM to classify the country that a question-answer pair is about.
    We will frame it as a task of: I have friends in several countries and I want to get the best answer possible. Who should I ask about this question? Choose from: English, French, German, Spanish, Russian.
    Question: [the question]
    Answer: [the answer] <---- this is where we probe the logits of the different languages and generate the weights.
    """

    languages_to_models_items = list(languages_to_models.items())
    languages_to_models_keys = [lang for lang, _ in languages_to_models_items]
    languages_to_models_values = [model for _, model in languages_to_models_items]

    # Tokenize the question and answer
    tokenized_text = instruct_model_tokenizer.apply_chat_template(
        conversation=[
            {
                "role": "system",
                f"content": f'You are a helpful assistant that identifies the topic of a question. The user will provide a question and an answer in JSON format, and you will choose the topic of both the question and the answer. The language of question and the answer doesn\'t matter, what matters is the topic. If the question and the answer are about mathematics, choose to ask "Math". If none of the contries is culturally related to the topic, choose to ask "Wikipedia" for encyclopedic knowledge. Write your response in JSON format. Follow this structure: {{"topic": "Topic: one of {json.dumps(languages_to_models_keys)}"}}.',
                # f"content": f'You are a helpful assistant that identifies the topic of a question. The user will provide a question and an answer in JSON format, and you will choose the topic of both the question and the answer. The language of question and the answer doesn\'t matter, what matters is the topic. Write your response in JSON format. Follow this structure: {{"topic": "Topic: one of {json.dumps(languages_to_models_keys)}"}}.',
            },
            {
                "role": "user",
                "content": json.dumps(
                    {
                        "question": "Where can I taste the best cassoulet in France?",
                        "answer": "Cassoulet is a typical dish of the Languedoc region, in the south of France, so I would recommend going there.",
                    },
                    **json_kwargs,
                ),
            },
            {
                "role": "assistant",
                "content": json.dumps({"topic": "Topic: French"}, **json_kwargs),
            },
            {
                "role": "user",
                "content": json.dumps(
                    {
                        "question": "Кто написал кодовую базу Windows 98?",
                        "answer": "Кодовая база Windows 98 была написана командой программистов-экспертов корпорации Microsoft.",
                    },
                    **json_kwargs,
                ),
            },
            {
                "role": "assistant",
                "content": json.dumps({"topic": "Topic: Wikipedia"}, **json_kwargs),
            },
            {
                "role": "user",
                "content": json.dumps(
                    {
                        "question": "Qué es mejor, la tortilla con cebolla o sin cebolla?",
                        "answer": "La tortilla con cebolla es mejor.",
                    },
                    **json_kwargs,
                ),
            },
            {
                "role": "assistant",
                "content": json.dumps({"topic": "Topic: Spanish"}, **json_kwargs),
            },
            {
                "role": "user",
                "content": json.dumps(
                    {
                        "question": "Tenho quantos anos se nasci em 1990 e hoje é 2022?",
                        "answer": "Se você nasceu em 1990 e hoje é 2022, você tem 32 anos.",
                    },
                    **json_kwargs,
                ),
            },
            {
                "role": "assistant",
                "content": json.dumps({"topic": "Topic: Math"}, **json_kwargs),
            },
            {
                "role": "user",
                "content": json.dumps(
                    {
                        "question": "Mitä meksikolaiseen tacoon ei saa koskaan laittaa?",
                        "answer": "Meksikolaiseen tacoon ei saa koskaan laittaa ananasta.",
                    },
                    **json_kwargs,
                ),
            },
            {
                "role": "assistant",
                "content": json.dumps({"topic": "Topic: Spanish"}, **json_kwargs),
            },
            {
                "role": "user",
                "content": json.dumps(
                    {
                        "question": "Chi fu il primo Cancelliere della Germania?",
                        "answer": "Il primo Cancelliere della Germania fu Konrad Adenauer.",
                    },
                    **json_kwargs,
                ),
            },
            {
                "role": "assistant",
                "content": json.dumps({"topic": "Topic: German"}, **json_kwargs),
            },
            {
                "role": "user",
                "content": json.dumps(
                    {"question": question, "answer": answer}, **json_kwargs
                ),
            },
            {
                "role": "assistant",
                "content": '{"topic": "Topic:',
            },
        ],
        add_generation_prompt=False,
        return_dict=True,
        return_tensors="pt",
        continue_final_message=True,
    )

    # # Remove the last token, which is the special token
    # tokenized_text["input_ids"] = tokenized_text["input_ids"][:, :-1]
    # tokenized_text["attention_mask"] = tokenized_text["attention_mask"][:, :-1]

    # Decode the tokenized text
    logging.debug(
        f"Last 20 tokens: {instruct_model_tokenizer.decode(tokenized_text['input_ids'][0, -20:])}"
    )

    # Run the model
    instruct_model_outputs = instruct_model(
        tokenized_text["input_ids"].to(instruct_model.device),
        attention_mask=tokenized_text["attention_mask"].to(instruct_model.device),
    )

    # Get the logits
    logits = instruct_model_outputs.logits.cpu()

    logging.debug(
        f"Top 20 next tokens: {instruct_model_tokenizer.convert_ids_to_tokens(logits[0, -1].topk(20).indices.tolist())} ({logits[0, -1].topk(20).values.tolist()})"
    )

    token_ids_to_select = instruct_model_tokenizer(
        [f" {lang}" for lang in languages_to_models_keys],
        add_special_tokens=False,
        return_tensors="pt",
        return_attention_mask=False,
        padding=True,
        padding_side="right",
    )

    # Take the first token of each language
    logging.debug(f"Token IDs: {(token_ids_to_select['input_ids'])}")
    token_ids_to_select = token_ids_to_select["input_ids"][:, 0]
    logging.debug(
        f"Token IDs to select: {instruct_model_tokenizer.convert_ids_to_tokens(token_ids_to_select)}"
    )

    # Get the logits for the languages we're interested in
    language_logits = logits[0, -1, token_ids_to_select]

    logging.debug(zip(languages_to_models_keys, language_logits.squeeze().tolist()))

    # Get the weights as the softmax of the logits
    language_logits = language_logits / weight_temperature
    weights = F.softmax(language_logits, dim=-1)

    weights_dict = {
        model_name: weight.item()
        for model_name, weight in zip(languages_to_models_values, weights.squeeze())
    }

    logging.debug(f"Weights: {weights_dict}")

    # Return the weights
    return weights_dict


@torch.no_grad()
def translate_text(
    text: str,
    translate_to: str,
    instruct_model: PreTrainedModel,
    instruct_model_tokenizer: PreTrainedTokenizer,
):
    """
    Translate a question-answer pair to a different language.
    """
    # Tokenize the question and answer
    tokenized_text = instruct_model_tokenizer.apply_chat_template(
        conversation=[
            {
                "role": "system",
                "content": f"You are an extremely precise translator that does not make any judgement or interpretation and only translates the text as it is written. Respect the exact meaning of the original text in your translation, write the original figures and names in the original language. Just write the translation directly, do not include any special markers, characters, formatting, notes or explanations, even if the information is incorrect.",
            },
            {
                "role": "user",
                "content": f"Translate to English: Der erste Mensch betrat den Mond 1951.",
            },
            {
                "role": "assistant",
                "content": "Man first stepped on the moon in 1951.",
            },
            {
                "role": "user",
                "content": f"Translate to English: Il nome del primo presidente degli Stati Uniti è Barack Obama, nato il 1° aprile 1870.",
            },
            {
                "role": "assistant",
                "content": "The name of the first president of the United States is Barack Obama, born on April 1, 1870.",
            },
            {
                "role": "user",
                "content": f"Translate to {translate_to}: {text}",
            },
        ],
        add_generation_prompt=True,
        return_dict=True,
        return_tensors="pt",
    )

    # Run the generation
    instruct_model_outputs = instruct_model.generate(
        tokenized_text["input_ids"].to(instruct_model.device),
        attention_mask=tokenized_text["attention_mask"].to(instruct_model.device),
        max_new_tokens=256,
        num_beams=3,
        early_stopping=True,
        no_repeat_ngram_size=3,
        do_sample=True,
        top_p=0.95,
        temperature=0.9,
        num_return_sequences=1,
    )

    logging.debug(
        f"Prompt: {instruct_model_tokenizer.decode(tokenized_text['input_ids'][0])}"
    )

    # Remove the part of the tokenized text that is the prompt
    instruct_model_outputs = instruct_model_outputs[
        ..., tokenized_text["input_ids"].size(-1) :
    ].squeeze()

    logging.debug(f"Instruct model outputs: {instruct_model_outputs}")

    # Decode the tokenized text
    translated_text = instruct_model_tokenizer.decode(
        instruct_model_outputs, skip_special_tokens=True
    )

    logging.debug(f"Translated text: {translated_text}")

    return translated_text


@torch.no_grad()
def translate_qa_pair(
    question: str,
    answer: str,
    translate_to: str,
    instruct_model: PreTrainedModel,
    instruct_model_tokenizer: PreTrainedTokenizer,
):
    """
    Translate a question-answer pair to a different language.
    """

    return {
        "translated_question": translate_text(
            text=question,
            translate_to=translate_to,
            instruct_model=instruct_model,
            instruct_model_tokenizer=instruct_model_tokenizer,
        ),
        "translated_answer": translate_text(
            text=answer,
            translate_to=translate_to,
            instruct_model=instruct_model,
            instruct_model_tokenizer=instruct_model_tokenizer,
        ),
    }

    # Tokenize the question and answer
    tokenized_text = instruct_model_tokenizer.apply_chat_template(
        conversation=[
            {
                "role": "system",
                f"content": f'You are a helpful assistant that translates a question-answer pair to a different language. The user will give a question and an answer in JSON format ({{"original_question": "original question", "original_answer": "original answer"}}), and you will have to translate it to {translate_to}. Write your response in JSON format ({{"translated_question": "translated question", "translated_answer": "translated answer"}}). Use double quotes for the strings.',
            },
            {
                "role": "user",
                "content": json.dumps(
                    {
                        "original_question": "Dove si trova la Torre di Pisa?",
                        "original_answer": "La Torre di Pisa si trova a Pisa.",
                    },
                    ensure_ascii=False,
                    indent=4,
                ),
            },
            {
                "role": "assistant",
                "content": json.dumps(
                    {
                        # "translated_language": translate_to,
                        "translated_question": "Where is the Tower of Pisa?",
                        "translated_answer": "The Tower of Pisa is in Pisa.",
                    },
                    ensure_ascii=False,
                    indent=4,
                ),
            },
            {
                "role": "user",
                "content": json.dumps(
                    {"original_question": question, "original_answer": answer},
                    ensure_ascii=False,
                    indent=4,
                ),
            },
        ],
        add_generation_prompt=True,
        return_dict=True,
        return_tensors="pt",
    )

    # Run the generation
    instruct_model_outputs = instruct_model.generate(
        tokenized_text["input_ids"].to(instruct_model.device),
        attention_mask=tokenized_text["attention_mask"].to(instruct_model.device),
        max_new_tokens=256,
        num_beams=4,
        early_stopping=True,
        no_repeat_ngram_size=3,
        do_sample=True,
        # top_k=50,
        # top_p=0.95,
        # temperature=1.0,
        num_return_sequences=1,
    )

    logging.info(
        f"Prompt: {instruct_model_tokenizer.decode(tokenized_text['input_ids'][0])}"
    )

    # Remove the part of the tokenized text that is the prompt
    instruct_model_outputs = instruct_model_outputs[
        ..., tokenized_text["input_ids"].size(-1) :
    ].squeeze()

    logging.debug(f"Instruct model outputs: {instruct_model_outputs}")

    # Decode the tokenized text
    translated_text = instruct_model_tokenizer.decode(
        instruct_model_outputs, skip_special_tokens=True
    )

    # Parse the translated text
    try:
        translated_text = json.loads(translated_text[0])
    except json.JSONDecodeError:
        logging.exception(f"Error decoding JSON: {translated_text}")
        # Use the original text
        translated_text = {"translated_question": question, "translated_answer": answer}

    return translated_text