# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/05_model.ipynb.

# %% auto 0
__all__ = ['ToolFormer']

# %% ../nbs/05_model.ipynb 4
from typing import Optional, List

import torch
from torch import nn
import torch.nn.functional as F

from transformers import AutoModelForCausalLM, AutoTokenizer
from torchtyping import TensorType

from .api import BaseAPI
from .utils import extract_api_request_content

# %% ../nbs/05_model.ipynb 6
class ToolFormer(nn.Module):
    def __init__(
        self,
        model: AutoModelForCausalLM,
        apis: List[BaseAPI],
        config: dict
    ):
        super().__init__()
        self.model = model
        self.apis = apis
        self.config = config
        self.is_calling_api: bool = False
        
        # TODO: make a config class contains token_id
        tokenizer = AutoTokenizer.from_pretrained(self.config["tokenizer"]["path"])
        self.tokenizer = tokenizer # TODO: remove after debug
        
        start_character = config["data_generator"]["api_start_character"]
        end_character = config["data_generator"]["api_end_character"]
        output_character = config["data_generator"]["api_output_character"]
        
        self.api_start_token_id = tokenizer(f' {start_character}', return_tensors="pt")["input_ids"][0]
        self.api_end_token_id = tokenizer(end_character, return_tensors="pt")["input_ids"][0]
        self.api_output_token_id = tokenizer(f'{output_character}', return_tensors="pt")["input_ids"][0]
    
        # TODO: support batch
        self.api_request_content: torch.Tensor = torch.tensor([])
    
    def _sampling(self, probs: TensorType["batch_size", "seq_len"]) -> TensorType["batch_size", "seq_len"]:
        return torch.argmax(probs, dim=-1)
    
    def execute_api(self, text_ids: TensorType["seq_len"]) -> TensorType["seq_len"]:
        """Execute an API call."""
        # content_ids = extract_api_request_content(text_ids, self.apis)
        pass
    
    def add_idx_to_api_request_content(self, idx: TensorType[1]):
        self.api_request_content = torch.cat([self.api_request_content, idx.unsqueeze(0)], dim=0)
    
    def forward(
        self,
        input_ids: TensorType["batch_size", "seq_len"],
        attention_mask: Optional[TensorType["batch_size", "seq_len"]]=None,
        max_new_tokens: int = 10,
        **kwargs
    ) -> TensorType["batch_size", "seq_len"]:
        # check padding to the left
        
        generated_ids = input_ids
        
        for _ in range(max_new_tokens):
            output_ids = self.model(
                input_ids=generated_ids,
                attention_mask=attention_mask,
                **kwargs
            )
            
            logits = output_ids.logits[:, -1, :]
            probs = F.softmax(logits, dim=-1)
            _, top_k_idx = torch.topk(probs, k=5, dim=-1)
            
            if self.is_calling_api is True:
                if self.api_end_token_id in top_k_idx:
                    # if the api end token is in the top_k_idx, then we will execute the api
                    # and then add api_end_token_id to the generated_ids
                    self.add_idx_to_api_request_content(self.api_end_token_id)
                    api_output_ids = self.execute_api(self.api_request_content)
                    pred_ids = torch.tensor([self.api_end_token_id, api_output_ids])
                    self.is_calling_api = False
                else:
                    pred_ids = self._sampling(probs)
                    self.add_idx_to_api_request_content(pred_ids)
            else:
                if self.api_start_token_id in top_k_idx:
                    # if the api start token is in the top_k_idx, then we are calling an api
                    self.is_calling_api = True
                    pred_ids = self.api_start_token_id
                    self.add_idx_to_api_request_content(pred_ids)
                else:
                    pred_ids = self._sampling(probs)
            
            generated_ids = torch.cat([generated_ids, pred_ids.unsqueeze(dim=1)], dim=1)
            attention_mask = torch.cat([attention_mask, torch.ones_like(pred_ids).unsqueeze(dim=1)], dim=1)
        
        return generated_ids
