README.md

@@ -1,11 +1,7 @@
 ---
 tags:
-- image-to-text
+- image-classification
-widget:
+library_name: generic
-- src: https://huggingface.co/datasets/mishig/sample_images/resolve/main/football-match.jpg
-  example_title: Football Match
-- src: https://huggingface.co/datasets/mishig/sample_images/resolve/main/dog-cat.jpg
-  example_title: Dog & Cat
 ---
 
 ## Example
@@ -16,89 +12,35 @@ as a proof-of-concept for the 🤗 FlaxVisionEncoderDecoder Framework.
 
 The model can be used as follows:
 
-**In PyTorch**
 ```python
 
-import torch
-import requests
-from PIL import Image
-from transformers import ViTFeatureExtractor, AutoTokenizer, VisionEncoderDecoderModel
-
-
-loc = "ydshieh/vit-gpt2-coco-en"
-
-feature_extractor = ViTFeatureExtractor.from_pretrained(loc)
-tokenizer = AutoTokenizer.from_pretrained(loc)
-model = VisionEncoderDecoderModel.from_pretrained(loc)
-model.eval()
-
-
-def predict(image):
-
-    pixel_values = feature_extractor(images=image, return_tensors="pt").pixel_values
-
-    with torch.no_grad():
-        output_ids = model.generate(pixel_values, max_length=16, num_beams=4, return_dict_in_generate=True).sequences
-
-    preds = tokenizer.batch_decode(output_ids, skip_special_tokens=True)
-    preds = [pred.strip() for pred in preds]
-
-    return preds
-
-
-# We will verify our results on an image of cute cats
-url = "http://images.cocodataset.org/val2017/000000039769.jpg"
-with Image.open(requests.get(url, stream=True).raw) as image:
-    preds = predict(image)
-
-print(preds)
-# should produce
-# ['a cat laying on top of a couch next to another cat']
-
-```
-
-**In Flax**
-```python
-
-import jax
 import requests
 from PIL import Image
 from transformers import ViTFeatureExtractor, AutoTokenizer, FlaxVisionEncoderDecoderModel
 
-
 loc = "ydshieh/vit-gpt2-coco-en"
 
 feature_extractor = ViTFeatureExtractor.from_pretrained(loc)
 tokenizer = AutoTokenizer.from_pretrained(loc)
 model = FlaxVisionEncoderDecoderModel.from_pretrained(loc)
 
-gen_kwargs = {"max_length": 16, "num_beams": 4}
-
-
-# This takes sometime when compiling the first time, but the subsequent inference will be much faster
-@jax.jit
-def generate(pixel_values):
-    output_ids = model.generate(pixel_values, **gen_kwargs).sequences
-    return output_ids
-    
-    
-def predict(image):
-
-    pixel_values = feature_extractor(images=image, return_tensors="np").pixel_values
-    output_ids = generate(pixel_values)
-    preds = tokenizer.batch_decode(output_ids, skip_special_tokens=True)
-    preds = [pred.strip() for pred in preds]
-    
-    return preds
-    
-    
 # We will verify our results on an image of cute cats
 url = "http://images.cocodataset.org/val2017/000000039769.jpg"
-with Image.open(requests.get(url, stream=True).raw) as image:
+with Image.open(requests.get(url, stream=True).raw) as img:
-    preds = predict(image)
+    pixel_values = feature_extractor(images=img, return_tensors="np").pixel_values
-    
+
+def generate_step(pixel_values):
+
+    output_ids = model.generate(pixel_values, max_length=16, num_beams=4).sequences
+    preds = tokenizer.batch_decode(output_ids, skip_special_tokens=True)
+    preds = [pred.strip() for pred in preds]
+
+    return preds
+
+preds = generate_step(pixel_values)
 print(preds)
+
 # should produce
 # ['a cat laying on top of a couch next to another cat']
 
-```
+```

pipeline.py

@@ -20,7 +20,7 @@ class PreTrainedPipeline():
         max_length = 16
         num_beams = 4
         # self.gen_kwargs = {"max_length": max_length, "num_beams": num_beams}
-        self.gen_kwargs = {"max_length": max_length, "num_beams": num_beams, "return_dict_in_generate": True, "output_scores": True}
+        self.gen_kwargs = {"max_length": max_length, "num_beams": num_beams, "return_dict_in_generate": True}
 
         self.model.to("cpu")
         self.model.eval()
@@ -29,12 +29,8 @@ class PreTrainedPipeline():
         def _generate(pixel_values):
 
             with torch.no_grad():
-                
+                output_ids = self.model.generate(pixel_values, **self.gen_kwargs).sequences
-                outputs = self.model.generate(pixel_values, **self.gen_kwargs)
+            return output_ids
-                output_ids = outputs.sequences
-                sequences_scores = outputs.sequences_scores
-                
-            return output_ids, sequences_scores
 
         self.generate = _generate
 
@@ -53,10 +49,10 @@ class PreTrainedPipeline():
         # pixel_values = self.feature_extractor(images=inputs, return_tensors="np").pixel_values
         pixel_values = self.feature_extractor(images=inputs, return_tensors="pt").pixel_values
 
-        output_ids, sequences_scores = self.generate(pixel_values)
+        output_ids = self.generate(pixel_values)
         preds = self.tokenizer.batch_decode(output_ids, skip_special_tokens=True)
         preds = [pred.strip() for pred in preds]
 
-        preds = [{"label": preds[0], "score": float(sequences_scores[0])}]
+        preds = [{"label": preds[0], "score": 1.0}]
 
         return preds

preprocessor_config.json

@@ -1,7 +1,6 @@
 {
   "do_normalize": true,
   "do_resize": true,
-  "feature_extractor_type": "ViTFeatureExtractor",
   "image_mean": [
     0.5,
     0.5,