import torch
from transformers import AutoTokenizer, VisionEncoderDecoderModel

import utils

class Inference:
  def __init__(self, decoder_model_name, max_length=32):
    self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

    self.tokenizer = AutoTokenizer.from_pretrained(decoder_model_name)
    self.encoder_decoder_model = VisionEncoderDecoderModel.from_pretrained('armgabrielyan/video-summarization')
    self.encoder_decoder_model.to(self.device)

    self.max_length = max_length

  def generate_text(self, video, encoder_model_name):
    if isinstance(video, str):
      pixel_values = utils.video2image_from_path(video, encoder_model_name)
    else:
      pixel_values = video

    if not self.tokenizer.pad_token:
      self.tokenizer.add_special_tokens({'pad_token': '[PAD]'})
      self.encoder_decoder_model.decoder.resize_token_embeddings(len(self.tokenizer))

    generated_ids = self.encoder_decoder_model.generate(pixel_values.unsqueeze(0).to(self.device), max_length=self.max_length)
    generated_text = self.tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]

    return generated_text