khang119966 commited on
Commit
1d87583
·
verified ·
1 Parent(s): a8325bc

Update README.md

Browse files
Files changed (1) hide show
  1. README.md +103 -0
README.md CHANGED
@@ -70,3 +70,106 @@ Dataset: [ViDoRe Benchmark](https://huggingface.co/collections/vidore/vidore-ben
70
  | ColVintern-1B | 0.9B | 78.8 | 71.6 | 48.3 | 84.6 | 92.9 | 88.7 | 89.4 | 95.2 | 59.6 |
71
  | Vintern-Embedding-1B | 0.9B | 82.85 | 75.37 | 51.79 | 86.2 | 97.52 | 93.19 | 93.97 | 97.09 | 67.72 |
72
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
70
  | ColVintern-1B | 0.9B | 78.8 | 71.6 | 48.3 | 84.6 | 92.9 | 88.7 | 89.4 | 95.2 | 59.6 |
71
  | Vintern-Embedding-1B | 0.9B | 82.85 | 75.37 | 51.79 | 86.2 | 97.52 | 93.19 | 93.97 | 97.09 | 67.72 |
72
 
73
+ ## Quickstart:
74
+
75
+ ```python
76
+ import torch
77
+ from PIL import Image
78
+ from transformers import AutoModel, AutoProcessor
79
+ import matplotlib.pyplot as plt
80
+
81
+ # ==============================
82
+ # 1. Load Model and Processor
83
+ # ==============================
84
+ model_name = "5CD-AI/Vintern-Embedding-1B"
85
+
86
+ processor = AutoProcessor.from_pretrained(
87
+ model_name,
88
+ trust_remote_code=True
89
+ )
90
+
91
+ model = AutoModel.from_pretrained(
92
+ model_name,
93
+ torch_dtype=torch.bfloat16, # Use bfloat16 for efficiency
94
+ low_cpu_mem_usage=True,
95
+ trust_remote_code=True,
96
+ ).eval().cuda() # Set model to eval mode and move to GPU
97
+
98
+ # ==============================
99
+ # 2. Prepare Input Data
100
+ # ==============================
101
+
102
+ # Images
103
+ images = [Image.open("ex1.jpg"), Image.open("ex2.jpg")]
104
+ batch_images = processor.process_images(images)
105
+
106
+ # Queries (questions)
107
+ queries = [
108
+ "Cảng Hải Phòng ở đâu ?",
109
+ "Phí giao hàng bao nhiêu ?",
110
+ ]
111
+ batch_queries = processor.process_queries(queries)
112
+
113
+ # Text documents
114
+ text_documents = [
115
+ "Cảng Hải Phòng là một cụm cảng biển tổng hợp cấp quốc gia, lớn thứ 2 ở Việt Nam sau cảng Sài Gòn, là cửa ngõ quốc tế của Việt Nam, nằm tại ba quận Hồng Bàng, Ngô Quyền và Hải An. Bên cạnh đó, cùng tên Cảng Hải Phòng (tiếng Anh: Port of Hai Phong hoặc Hai Phong Port) là một cụm cảng biển thuộc Công ty cổ phần cảng Hải Phòng tại thành phố Hải Phòng, Việt Nam. Đây là một trong hai cảng biển tổng hợp lớn và lâu đời nhất tại Việt Nam, cùng với Công ty Cảng Sài Gòn ở phía Nam.",
116
+ "Sân bay Chu Lai (tỉnh Quảng Nam) cũng được hãng hàng không giá rẻ Vietjet đề xuất đầu tư nâng cấp 20.000 tỉ đồng theo 3 giai đoạn từ 2020-2025 để đến năm 2025 trở thành Cảng hàng không quốc tế và trở thành trung tâm trung chuyển, vận tải hàng hóa lớn của cả nước theo quy hoạch của Bộ GTVT năm 2015.",
117
+ ]
118
+ batch_text_docs = processor.process_docs(text_documents)
119
+
120
+ # Combine images and text docs for retrieval
121
+ raw_docs = images + text_documents
122
+
123
+ # ==============================
124
+ # 3. Move Tensors to GPU
125
+ # ==============================
126
+ # Images
127
+ batch_images["pixel_values"] = batch_images["pixel_values"].cuda().bfloat16()
128
+ batch_images["input_ids"] = batch_images["input_ids"].cuda()
129
+ batch_images["attention_mask"] = batch_images["attention_mask"].cuda().bfloat16()
130
+
131
+ # Queries
132
+ batch_queries["input_ids"] = batch_queries["input_ids"].cuda()
133
+ batch_queries["attention_mask"] = batch_queries["attention_mask"].cuda().bfloat16()
134
+
135
+ # Text Documents
136
+ batch_text_docs["input_ids"] = batch_text_docs["input_ids"].cuda()
137
+ batch_text_docs["attention_mask"] = batch_text_docs["attention_mask"].cuda().bfloat16()
138
+
139
+ # ==============================
140
+ # 4. Generate Embeddings
141
+ # ==============================
142
+ with torch.no_grad():
143
+ image_embeddings = model(**batch_images)
144
+ query_embeddings = model(**batch_queries)
145
+ text_docs_embeddings = model(**batch_text_docs)
146
+
147
+ # ==============================
148
+ # 5. Compute Similarity Scores
149
+ # ==============================
150
+ scores = processor.score_multi_vector(
151
+ query_embeddings,
152
+ list(image_embeddings) + list(text_docs_embeddings)
153
+ )
154
+
155
+ max_scores, max_indices = torch.max(scores, dim=1)
156
+
157
+ # ==============================
158
+ # 6. Print Results
159
+ # ==============================
160
+ for i, query in enumerate(queries):
161
+ print("=" * 100)
162
+ print(f"Query: '{query}'")
163
+ print(f"Score: {max_scores[i].item()}\n")
164
+
165
+ doc = raw_docs[max_indices[i]]
166
+ if isinstance(doc, str):
167
+ print(f"Matched Text Document:\n{doc}\n")
168
+ else:
169
+ plt.figure(figsize=(5, 5))
170
+ plt.imshow(doc)
171
+ plt.axis("off")
172
+ plt.show()
173
+
174
+ ```
175
+