add batch generation

app.py

@@ -2,6 +2,7 @@ from datetime import timedelta
 import gradio as gr
 from sentence_transformers import SentenceTransformer
 import torchvision
+import torch
 from sklearn.metrics.pairwise import cosine_similarity
 import numpy as np
 
@@ -27,13 +28,10 @@ def search_in_video(video, query):
         video[idx:idx + frame_step, :, :, :] for idx in range(0, video.shape[0], frame_step)
     ]
 
-    generated_texts = []
+    pixel_values = [utils.video2image(video_seg, encoder_model_name) for video_seg in video_segments]
+    pixel_values = torch.stack(pixel_values)
 
-    for video_seg in video_segments:
-        pixel_values = utils.video2image(video_seg, encoder_model_name)
-
-        generated_text = inference.generate_text(pixel_values, encoder_model_name)
-        generated_texts.append(generated_text)
+    generated_texts = inference.generate_texts(pixel_values)
 
     sentences = [query] + generated_texts
 

inference.py

@@ -1,7 +1,6 @@
 import torch
 from transformers import AutoTokenizer, VisionEncoderDecoderModel
 
-import utils
 
 class Inference:
   def __init__(self, decoder_model_name, max_length=32):
@@ -13,22 +12,17 @@ class Inference:
 
     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
-
+  def generate_texts(self, pixel_values):
     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),
+      pixel_values.to(self.device),
       max_length=self.max_length,
       num_beams=4, 
       no_repeat_ngram_size=2,
     )
-    generated_text = self.tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
+    generated_texts = self.tokenizer.batch_decode(generated_ids, skip_special_tokens=True)
 
-    return generated_text
+    return generated_texts

utils.py

@@ -1,15 +1,8 @@
 from transformers import ViTFeatureExtractor
-import torchvision
 import torchvision.transforms.functional as fn
 import torch as th
 
 
-def video2image_from_path(video_path, feature_extractor_name):
-    video = torchvision.io.read_video(video_path)
-
-    return video2image(video[0], feature_extractor_name)
-
-
 def video2image(video, feature_extractor_name):
     feature_extractor = ViTFeatureExtractor.from_pretrained(
         feature_extractor_name