from shared.utils.phase_progress import text_encoding_progress
from torch import Tensor, nn
from transformers import CLIPTextModel, CLIPTokenizer, T5EncoderModel, T5Tokenizer
import os

class HFEmbedder(nn.Module):
    def __init__(self, version: str, text_encoder_filename, max_length: int, is_clip = False, **hf_kwargs):
        super().__init__()
        self.is_clip = is_clip
        self.max_length = max_length
        self.output_key = "pooler_output" if self.is_clip else "last_hidden_state"

        if is_clip:
            from mmgp import offload
            self.tokenizer: CLIPTokenizer = CLIPTokenizer.from_pretrained(version, max_length=max_length)
            self.hf_module= offload.fast_load_transformers_model(os.path.join(version, "model.safetensors"), ignore_unused_weights= True, modelClass=CLIPTextModel, forcedConfigPath = os.path.join(version, "text_config.json"), writable_tensors=False)
            # self.model.final_layer_norm = self.model.text_model.final_layer_norm
            # self.hf_module: CLIPTextModel = CLIPTextModel.from_pretrained(version, **hf_kwargs)
        else:
            from mmgp import offload as offloadobj
            self.tokenizer: T5Tokenizer = T5Tokenizer.from_pretrained(os.path.dirname(text_encoder_filename),  max_length=max_length)
            self.hf_module: T5EncoderModel  = offloadobj.fast_load_transformers_model(text_encoder_filename, writable_tensors=False) 

        self.hf_module = self.hf_module.eval().requires_grad_(False)

    def forward(self, text: list[str]) -> Tensor:
        batch_encoding = self.tokenizer(
            text,
            truncation=True,
            max_length=self.max_length,
            return_length=False,
            return_overflowing_tokens=False,
            padding="max_length",
            return_tensors="pt",
        )

        with text_encoding_progress(self.hf_module.text_model.encoder.layers if self.is_clip else self.hf_module.encoder.block, prompt_count=len(text)):
            outputs = self.hf_module(
                input_ids=batch_encoding["input_ids"].to(self.hf_module.device),
                attention_mask=None,
                output_hidden_states=False,
            )
        return outputs[self.output_key].bfloat16()