Update README.md
Browse files
README.md
CHANGED
|
@@ -73,7 +73,40 @@ Users (both direct and downstream) should be made aware of the risks, biases and
|
|
| 73 |
|
| 74 |
Use the code below to get started with the model.
|
| 75 |
|
| 76 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 77 |
|
| 78 |
## Training Details
|
| 79 |
|
|
|
|
| 73 |
|
| 74 |
Use the code below to get started with the model.
|
| 75 |
|
| 76 |
+
|
| 77 |
+
```python
|
| 78 |
+
from PIL import Image
|
| 79 |
+
from transformers import AutoProcessor, AutoModelForCausalLM
|
| 80 |
+
|
| 81 |
+
device = "cuda:0" if torch.cuda.is_available() else "cpu"
|
| 82 |
+
torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
|
| 83 |
+
|
| 84 |
+
model = AutoModelForCausalLM.from_pretrained("AskUI/PTA-1", torch_dtype=torch_dtype, trust_remote_code=True).to(device)
|
| 85 |
+
processor = AutoProcessor.from_pretrained("AskUI/PTA-1", trust_remote_code=True)
|
| 86 |
+
|
| 87 |
+
task_prompt = "<OPEN_VOCABULARY_DETECTION>"
|
| 88 |
+
prompt = task_prompt + "description of the target element"
|
| 89 |
+
|
| 90 |
+
image = Image.open("path to screenshot")
|
| 91 |
+
|
| 92 |
+
inputs = processor(text=prompt, images=image, return_tensors="pt").to(device, torch_dtype)
|
| 93 |
+
|
| 94 |
+
generated_ids = model.generate(
|
| 95 |
+
input_ids=inputs["input_ids"],
|
| 96 |
+
pixel_values=inputs["pixel_values"],
|
| 97 |
+
max_new_tokens=1024,
|
| 98 |
+
do_sample=False,
|
| 99 |
+
num_beams=3,
|
| 100 |
+
)
|
| 101 |
+
generated_text = processor.batch_decode(generated_ids, skip_special_tokens=False)[0]
|
| 102 |
+
|
| 103 |
+
parsed_answer = processor.post_process_generation(generated_text, task="<OPEN_VOCABULARY_DETECTION>", image_size=(image.width, image.height))
|
| 104 |
+
|
| 105 |
+
print(parsed_answer)
|
| 106 |
+
```
|
| 107 |
+
|
| 108 |
+
|
| 109 |
+
|
| 110 |
|
| 111 |
## Training Details
|
| 112 |
|