Update README.md

README.md

@@ -8,97 +8,156 @@ library_name: transformers
 
 Please download and save `builder.py`, `conversation.py` locally.
 
-### Basic Text Generation
-```python
-import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM
-
-# load the model and tokenizer
-model_name = "jadechoghari/ferret-gemma"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16)
-
-# give input text
-input_text = "The United States of America is a country situated on earth"
-
-# tokenize the input text
-inputs = tokenizer(input_text, return_tensors="pt", padding=True).to("cuda" if torch.cuda.is_available() else "cpu")
-
-model = model.to("cuda" if torch.cuda.is_available() else "cpu")
-
-output = model.generate(inputs['input_ids'], max_length=50, num_return_sequences=1)
-
-# decode and print the output
-generated_text = tokenizer.decode(output[0], skip_special_tokens=True)
-print(generated_text)
-```
-
-### Image and Text Generation
 ```python
 import torch
 from PIL import Image
 from conversation import conv_templates
-from builder import load_pretrained_model  # custom model loader
+from builder import load_pretrained_model  # Assuming this is your custom model loader
+from functools import partial
+import numpy as np
+
+# define the task categories
+box_in_tasks = ['widgetcaptions', 'taperception', 'ocr', 'icon_recognition', 'widget_classification', 'example_0']
+box_out_tasks = ['widget_listing', 'find_text', 'find_icons', 'find_widget', 'conversation_interaction']
+no_box_tasks = ['screen2words', 'detailed_description', 'conversation_perception', 'gpt4']
+
+# function to generate the mask
+def generate_mask_for_feature(coor, raw_w, raw_h, mask=None):
+    """
+    Generates a region mask based on provided coordinates.
+    Handles both point and box input.
+    """
+    if mask is not None:
+        assert mask.shape[0] == raw_w and mask.shape[1] == raw_h
+    coor_mask = np.zeros((raw_w, raw_h))
+
+    # if it's a point (2 coordinates)
+    if len(coor) == 2:
+        span = 5  # Define the span for the point
+        x_min = max(0, coor[0] - span)
+        x_max = min(raw_w, coor[0] + span + 1)
+        y_min = max(0, coor[1] - span)
+        y_max = min(raw_h, coor[1] + span + 1)
+        coor_mask[int(x_min):int(x_max), int(y_min):int(y_max)] = 1
+        assert (coor_mask == 1).any(), f"coor: {coor}, raw_w: {raw_w}, raw_h: {raw_h}"
+
+    # if it's a box (4 coordinates)
+    elif len(coor) == 4:
+        coor_mask[coor[0]:coor[2]+1, coor[1]:coor[3]+1] = 1
+        if mask is not None:
+            coor_mask = coor_mask * mask
+
+    # Convert to torch tensor and ensure it contains non-zero values
+    coor_mask = torch.from_numpy(coor_mask)
+    assert len(coor_mask.nonzero()) != 0, "Generated mask is empty :("
 
-# load model and tokenizer, then preprocess an image
-def infer_single_prompt(image_path, prompt, model_path):
+    return coor_mask
+```
+### Now, define the infer function
+```python
+def infer_single_prompt(image_path, prompt, model_path, region=None, model_name="ferret_gemma", conv_mode="ferret_gemma_instruct"):
     img = Image.open(image_path).convert('RGB')
-    tokenizer, model, image_processor, _ = load_pretrained_model(model_path, None, "ferret_gemma")
-    image_tensor = image_processor.preprocess(img, return_tensors='pt', size=(336, 336))['pixel_values'][0].unsqueeze(0).half()
 
-    # prepare prompt
-    conv = conv_templates["ferret_gemma_instruct"].copy()
-    conv.append_message(conv.roles[0], f"Image and prompt: {prompt}")
-    input_ids = tokenizer(conv.get_prompt(), return_tensors='pt')['input_ids'].cuda()
-
-    image_tensor = image_tensor.cuda()
-
-    # generate text output
-    with torch.inference_mode():
-        output_ids = model.generate(input_ids, images=image_tensor, max_new_tokens=1024)
-    
-    # decode the output
-    output_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
-    return output_text.strip()
+    # this loads the model, image processor and tokenizer
+    tokenizer, model, image_processor, context_len = load_pretrained_model(model_path, None, model_name)
 
-# Usage
-result = infer_single_prompt("image.jpg", "Describe the contents of the image.", "jadechoghari/ferret-gemma")
-print(result)
-```
+    # define the image size (e.g., 224x224 or 336x336)
+    image_size = {"height": 336, "width": 336}
 
-### Text, Image, and Bounding Box
-```python
-import torch
-from PIL import Image
-from functools import partial
-from builder import load_pretrained_model
+    # process the image
+    image_tensor = image_processor.preprocess(
+        img,
+        return_tensors='pt',
+        do_resize=True,
+        do_center_crop=False,
+        size=(image_size['height'], image_size['width'])
+    )['pixel_values'][0].unsqueeze(0)
 
-# generates a bounding box mask
-def generate_mask_for_feature(coor, img_w, img_h):
-    coor_mask = torch.zeros((img_w, img_h))
-    coor_mask[coor[0]:coor[2]+1, coor[1]:coor[3]+1] = 1
-    return coor_mask
+    image_tensor = image_tensor.half().cuda()
 
-def infer_with_bounding_box(image_path, prompt, model_path, region):
-    img = Image.open(image_path).convert('RGB')
-    tokenizer, model, image_processor, _ = load_pretrained_model(model_path, None, "ferret_gemma")
-    image_tensor = image_processor.preprocess(img, return_tensors='pt', size=(336, 336))['pixel_values'][0].unsqueeze(0).half().cuda()
+    # generate the prompt per template requirement
+    conv = conv_templates[conv_mode].copy()
+    conv.append_message(conv.roles[0], prompt)
+    conv.append_message(conv.roles[1], None)
+    prompt_input = conv.get_prompt()
 
-    input_ids = tokenizer(f"Image and prompt: {prompt}", return_tensors='pt')['input_ids'].cuda()
+    # tokenize prompt
+    input_ids = tokenizer(prompt_input, return_tensors='pt')['input_ids'].cuda()
 
-    # create region mask
-    mask = generate_mask_for_feature(region, *img.size).unsqueeze(0).half().cuda()
+    # region mask logic (if region is provided)
+    region_masks = None
+    if region is not None:
+        raw_w, raw_h = img.size
+        region_masks = generate_mask_for_feature(region, raw_w, raw_h).unsqueeze(0).cuda().half()
+        region_masks = [[region_masks]]  # Wrap the mask in lists as expected by the model
 
-    # generate output with region mask
+    # generate model output
     with torch.inference_mode():
+        # Use region_masks in model's forward call
         model.orig_forward = model.forward
-        model.forward = partial(model.orig_forward, region_masks=[[mask]])
-        output_ids = model.generate(input_ids, images=image_tensor, max_new_tokens=1024)
-    
-    output_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
+        model.forward = partial(
+            model.orig_forward,
+            region_masks=region_masks
+        )
+        output_ids = model.generate(
+            input_ids,
+            images=image_tensor,
+            max_new_tokens=1024,
+            num_beams=1,
+            region_masks=region_masks,  # pass the region mask to the model
+            image_sizes=[img.size]
+        )
+        model.forward = model.orig_forward
+
+    # we decode the output
+    output_text = tokenizer.batch_decode(output_ids, skip_special_tokens=True)[0]
     return output_text.strip()
+```
 
-# Usage
-result = infer_with_bounding_box("image.jpg", "Describe the contents of the box.", "jadechoghari/ferret-gemma", (50, 50, 200, 200))
-print(result)
+# We also define a task-specific inference function
+```python
+def infer_ui_task(image_path, prompt, model_path, task, region=None):
+    """
+    Handles task types: box_in_tasks, box_out_tasks, no_box_tasks.
+    """
+    if task in box_in_tasks and region is None:
+        raise ValueError(f"Task {task} requires a bounding box region.")
+    
+    if task in box_in_tasks:
+        print(f"Processing {task} with bounding box region.")
+        return infer_single_prompt(image_path, prompt, model_path, region)
+    
+    elif task in box_out_tasks:
+        print(f"Processing {task} without bounding box region.")
+        return infer_single_prompt(image_path, prompt, model_path)
+    
+    elif task in no_box_tasks:
+        print(f"Processing {task} without image or bounding box.")
+        return infer_single_prompt(image_path, prompt, model_path)
+    
+    else:
+        raise ValueError(f"Unknown task type: {task}")
+```
+
+### Usage:
+```python
+# Example image and model paths
+image_path = 'image.jpg'
+model_path = 'jadechoghari/ferret-gemma'
+
+# Task requiring bounding box
+task = 'widgetcaptions'
+region = (50, 50, 200, 200)
+result = infer_ui_task(image_path, "Describe the contents of the box.", model_path, task, region=region)
+print("Result:", result)
+
+# Task not requiring bounding box
+task = 'conversation_interaction'
+result = infer_ui_task(image_path, "How do I navigate to the Games tab?", model_path, task)
+print("Result:", result)
+
+# Task with no image processing
+task = 'detailed_description'
+result = infer_ui_task(image_path, "Please describe the screen in detail.", model_path, task)
+print("Result:", result)
 ```