from typing import List, Tuple
import torch
import ollama

from lm_eval.api.registry import register_model

from src.backend.hflm_with_measurement import HFLMWithMeasurement

@register_model("ollama")
class OllamaChatTemplate(HFLMWithMeasurement):
    def __init__(self, use_chat_template=True, **kwargs):
        super().__init__(**kwargs)
        self.use_chat_template = use_chat_template
        # Initialize the ollama model and tokenizer here
        self.model = ollama.OllamaModel.from_pretrained(kwargs['model_name_or_path'])
        self.tokenizer = ollama.OllamaTokenizer.from_pretrained(kwargs['model_name_or_path'])

    def tok_batch_encode(
        self,
        strings: List[str],
        padding_side: str = "left",
        left_truncate_len: int = None,
        truncation: bool = False,
    ) -> Tuple[torch.Tensor, torch.Tensor]:

        if self.use_chat_template:
            try:
                updated_strings = []
                for input_string in strings:
                    messages = [
                        {"role": "user", "content": f"{input_string}"},
                    ]
                    updated_string = self.tokenizer.apply_chat_template(messages, tokenize=False)
                    updated_strings.append(updated_string)
                strings = updated_strings[:]
            except Exception as e:
                print(f"Failed to update input string with chat template: {e}")

        # Encode a batch of strings. Converts to tensors and pads automatically.
        old_padding_side = self.tokenizer.padding_side
        self.tokenizer.padding_side = padding_side

        encoding = self.tokenizer(
            strings,
            truncation=truncation,
            padding="longest",
            return_tensors="pt",
            add_special_tokens=True,
        )
        if left_truncate_len:
            encoding["input_ids"] = encoding["input_ids"][:, -left_truncate_len:]
            encoding["attention_mask"] = encoding["attention_mask"][:, -left_truncate_len:]
        self.tokenizer.padding_side = old_padding_side

        return encoding["input_ids"], encoding["attention_mask"]