Delete test.py
Browse files
test.py
DELETED
@@ -1,29 +0,0 @@
|
|
1 |
-
import torch
|
2 |
-
from transformers import AutoModelForCausalLM, AutoTokenizer
|
3 |
-
from PIL import Image
|
4 |
-
|
5 |
-
|
6 |
-
torch.set_default_device("cuda")
|
7 |
-
|
8 |
-
#Create model
|
9 |
-
model = AutoModelForCausalLM.from_pretrained(
|
10 |
-
"/data/ouyangxc/labs/hg/imp-v1-3b",
|
11 |
-
torch_dtype=torch.float16,
|
12 |
-
device_map="auto",
|
13 |
-
trust_remote_code=True)
|
14 |
-
tokenizer = AutoTokenizer.from_pretrained("/data/ouyangxc/labs/hg/imp-v1-3b", trust_remote_code=True)
|
15 |
-
|
16 |
-
#Set inputs
|
17 |
-
text = "A chat between a curious user and an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the user's questions. USER: <image>\nWhat's the color of the car? ASSISTANT:"
|
18 |
-
image = Image.open("images/car.jpg")
|
19 |
-
|
20 |
-
input_ids = tokenizer(text, return_tensors='pt').input_ids
|
21 |
-
image_tensor = model.image_preprocess(image)
|
22 |
-
|
23 |
-
#Generate the answer
|
24 |
-
output_ids = model.generate(
|
25 |
-
input_ids,
|
26 |
-
max_new_tokens=150,
|
27 |
-
images=image_tensor,
|
28 |
-
use_cache=True)[0]
|
29 |
-
print(tokenizer.decode(output_ids[input_ids.shape[1]:], skip_special_tokens=True).strip())
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|