Armen Gabrielyan commited on
Commit
16b0970
1 Parent(s): 68da745

change from auto model to vision encoder decoder model

Browse files
Files changed (3) hide show
  1. inference.py +2 -2
  2. requirements.txt +0 -2
  3. utils.py +0 -2
inference.py CHANGED
@@ -1,5 +1,5 @@
1
  import torch
2
- from transformers import AutoTokenizer, AutoModel
3
 
4
  import utils
5
 
@@ -8,7 +8,7 @@ class Inference:
8
  self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
9
 
10
  self.tokenizer = AutoTokenizer.from_pretrained(decoder_model_name)
11
- self.encoder_decoder_model = AutoModel.from_pretrained('armgabrielyan/video-summarization')
12
  self.encoder_decoder_model.to(self.device)
13
 
14
  self.max_length = max_length
 
1
  import torch
2
+ from transformers import AutoTokenizer, VisionEncoderDecoderModel
3
 
4
  import utils
5
 
 
8
  self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
9
 
10
  self.tokenizer = AutoTokenizer.from_pretrained(decoder_model_name)
11
+ self.encoder_decoder_model = VisionEncoderDecoderModel.from_pretrained('armgabrielyan/video-summarization')
12
  self.encoder_decoder_model.to(self.device)
13
 
14
  self.max_length = max_length
requirements.txt CHANGED
@@ -1,4 +1,2 @@
1
- nltk==3.7
2
- tqdm==4.64.0
3
  scikit-learn==1.1.1
4
  sentence-transformers==2.2.0
 
 
 
1
  scikit-learn==1.1.1
2
  sentence-transformers==2.2.0
utils.py CHANGED
@@ -2,8 +2,6 @@ from transformers import ViTFeatureExtractor
2
  import torchvision
3
  import torchvision.transforms.functional as fn
4
  import torch as th
5
- import os
6
- import pickle
7
 
8
 
9
  def video2image_from_path(video_path, feature_extractor_name):
 
2
  import torchvision
3
  import torchvision.transforms.functional as fn
4
  import torch as th
 
 
5
 
6
 
7
  def video2image_from_path(video_path, feature_extractor_name):