|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
import tqdm |
|
from PIL import Image |
|
import hashlib |
|
import torch |
|
import fitz |
|
|
|
|
|
def get_image_md5(img: Image.Image): |
|
img_byte_array = img.tobytes() |
|
hash_md5 = hashlib.md5() |
|
hash_md5.update(img_byte_array) |
|
hex_digest = hash_md5.hexdigest() |
|
return hex_digest |
|
|
|
def pdf_to_images(pdf_path, dpi=200): |
|
doc = fitz.open(pdf_path) |
|
images = [] |
|
for page in tqdm.tqdm(doc): |
|
pix = page.get_pixmap(dpi=dpi) |
|
img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples) |
|
images.append(img) |
|
return images |
|
|
|
|
|
class PDFVisualRetrieval: |
|
def __init__(self, model, tokenizer): |
|
self.tokenizer = tokenizer |
|
self.model = model |
|
self.reps = {} |
|
self.images = {} |
|
|
|
def add_visual_documents(self, knowledge_base_name: str, images: Image.Image): |
|
if knowledge_base_name not in self.reps: |
|
self.reps[knowledge_base_name] = {} |
|
if knowledge_base_name not in self.images: |
|
self.images[knowledge_base_name] = {} |
|
for image in tqdm.tqdm(images): |
|
image_md5 = get_image_md5(image) |
|
with torch.no_grad(): |
|
reps = self.model(text=[''], image=[image], tokenizer=self.tokenizer).reps |
|
self.reps[knowledge_base_name][image_md5] = reps.squeeze(0) |
|
self.images[knowledge_base_name][image_md5] = image |
|
return |
|
|
|
def retrieve(self, knowledge_base: str, query: str, topk: int): |
|
doc_reps = list(self.reps[knowledge_base].values()) |
|
query_with_instruction = "Represent this query for retrieving relavant document: " + query |
|
with torch.no_grad(): |
|
query_rep = self.model(text=[query_with_instruction], image=[None], tokenizer=self.tokenizer).reps.squeeze(0) |
|
doc_reps_cat = torch.stack(doc_reps, dim=0) |
|
similarities = torch.matmul(query_rep, doc_reps_cat.T) |
|
topk_values, topk_doc_ids = torch.topk(similarities, k=topk) |
|
topk_values_np = topk_values.cpu().numpy() |
|
topk_doc_ids_np = topk_doc_ids.cpu().numpy() |
|
similarities_np = similarities.cpu().numpy() |
|
all_images_doc_list = list(self.images[knowledge_base].values()) |
|
images_topk = [all_images_doc_list[idx] for idx in topk_doc_ids_np] |
|
return topk_doc_ids_np, topk_values_np, images_topk |
|
|
|
def add_pdf(self, knowledge_base_name: str, pdf_file_path: str, dpi: int = 200): |
|
print("[1/2] rendering pdf to images..") |
|
images = pdf_to_images(pdf_file_path, dpi=dpi) |
|
print("[2/2] model encoding images..") |
|
self.add_visual_documents(knowledge_base_name=knowledge_base_name, images=images) |
|
print("add pdf ok.") |
|
return |
|
|
|
|
|
if __name__ == "__main__": |
|
from transformers import AutoModel |
|
from transformers import AutoTokenizer |
|
from PIL import Image |
|
import torch |
|
|
|
device = 'cuda:0' |
|
|
|
|
|
model_path = '/home/jeeves/xubokai/minicpm-visual-embedding-v0' |
|
tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True) |
|
model = AutoModel.from_pretrained(model_path, trust_remote_code=True) |
|
model.to(device) |
|
|
|
pdf_path = "/home/jeeves/xubokai/minicpm-visual-embedding-v0/2406.07422v1.pdf" |
|
retriever = PDFVisualRetrieval(model=model, tokenizer=tokenizer) |
|
retriever.add_pdf('test', pdf_path) |
|
|
|
topk_doc_ids_np, topk_values_np, images_topk = retriever.retrieve(knowledge_base='test', query='what is the number of VQ of this kind of codec method?', topk=5) |
|
|
|
topk_doc_ids_np, topk_values_np, images_topk = retriever.retrieve(knowledge_base='test', query='the training loss curve of this paper?', topk=5) |
|
|
|
topk_doc_ids_np, topk_values_np, images_topk = retriever.retrieve(knowledge_base='test', query='the experiment table?', topk=5) |
|
|
|
|
|
|
