########################################################################################################
# The EmbeddingRWKV Model - https://github.com/add_later
########################################################################################################

from types import SimpleNamespace
from typing import List, Optional, Sequence, Tuple, Union

import torch
import torch.nn as nn

from reference.rwkv7 import RWKV_x070


Tokens = List[int]
BatchTokens = Sequence[Sequence[int]]


class EmbeddingRWKV(nn.Module):
    def __init__(self, model_path: str):
        super().__init__()

        model_prefix = model_path[:-4] if model_path.endswith(".pth") else model_path
        print(f"Loading {model_prefix}.pth \n")

        args = SimpleNamespace(MODEL_NAME=model_prefix)
        self.rwkv = RWKV_x070(args)

    def generate_zero_state(self, bsz):
        return self.rwkv.generate_zero_state(bsz)

    @torch.inference_mode()
    def forward(
        self,
        tokens: Union[Tokens, BatchTokens],
        state: Optional[List[torch.Tensor]] = None,
        full_output: bool = False,
    ):
        is_batch = isinstance(tokens, Sequence) and len(tokens) > 0 and isinstance(tokens[0], Sequence)
        if state is None:
            B = len(tokens) if is_batch else 0
            state = self.generate_zero_state(B)
        if is_batch:
            output = self.rwkv.forward_batch(tokens, state, full_output)
        else:
            output = self.rwkv.forward(tokens, state, full_output)
        return output, state
