kevin-pek commited on
Commit
e8c3179
·
1 Parent(s): ef09334

replace NougatProcessor with AutoProcessor

Browse files
Files changed (2) hide show
  1. config.json +1 -1
  2. handler.py +3 -3
config.json CHANGED
@@ -184,5 +184,5 @@
184
  "model_type": "vision-encoder-decoder",
185
  "tie_word_embeddings": false,
186
  "torch_dtype": "float32",
187
- "transformers_version": "main"
188
  }
 
184
  "model_type": "vision-encoder-decoder",
185
  "tie_word_embeddings": false,
186
  "torch_dtype": "float32",
187
+ "transformers_version": "4.34.0.dev0"
188
  }
handler.py CHANGED
@@ -1,13 +1,13 @@
1
  from io import BytesIO
2
- from typing import Dict, List, Any
3
- from transformers import NougatProcessor, VisionEncoderDecoderModel
4
  from transformers.image_utils import base64
5
  from PIL import Image
6
  import torch
7
 
8
  class EndpointHandler():
9
  def __init__(self, path="facebook/nougat-base") -> None:
10
- self.processor = NougatProcessor.from_pretrained(path)
11
  self.model = VisionEncoderDecoderModel.from_pretrained(path)
12
 
13
  self.device = "cuda" if torch.cuda.is_available() else "cpu"
 
1
  from io import BytesIO
2
+ from typing import Dict, Any
3
+ from transformers import AutoProcessor, VisionEncoderDecoderModel
4
  from transformers.image_utils import base64
5
  from PIL import Image
6
  import torch
7
 
8
  class EndpointHandler():
9
  def __init__(self, path="facebook/nougat-base") -> None:
10
+ self.processor = AutoProcessor.from_pretrained(path)
11
  self.model = VisionEncoderDecoderModel.from_pretrained(path)
12
 
13
  self.device = "cuda" if torch.cuda.is_available() else "cpu"