Update README.md

README.md

@@ -51,7 +51,8 @@ from transformers import AutoModel, AutoTokenizer
 import torch
 import torch.nn.functional as F
 from PIL import Image
-import os
+import requests
+from io import BytesIO
 
 def weighted_mean_pooling(hidden, attention_mask):
     attention_mask_ = attention_mask * attention_mask.cumsum(dim=1)
@@ -83,20 +84,25 @@ def encode(text_or_image_list):
     embeddings = F.normalize(reps, p=2, dim=1).detach().cpu().numpy()
     return embeddings
 
-tokenizer = AutoTokenizer.from_pretrained("openbmb/VisRAG-Ret", trust_remote_code=True)
-model = AutoModel.from_pretrained("openbmb/VisRAG-Ret", torch_dtype=torch.bfloat16, trust_remote_code=True)
+tokenizer = AutoTokenizer.from_pretrained("/mnt/data/user/tc_agi/klara/datasets/visrag_ret/visrag_ret", trust_remote_code=True)
+model = AutoModel.from_pretrained("/mnt/data/user/tc_agi/klara/datasets/visrag_ret/visrag_ret", torch_dtype=torch.bfloat16, trust_remote_code=True)
 model.eval()
 
-script_dir = os.path.dirname(os.path.realpath(__file__))
 queries = ["What does a dog look like?"]
-passages = [
-    Image.open(os.path.join(script_dir, 'test_image/cat.jpeg')).convert('RGB'),
-    Image.open(os.path.join(script_dir, 'test_image/dog.jpg')).convert('RGB'),
-]
-
 INSTRUCTION = "Represent this query for retrieving relevant documents: "
 queries = [INSTRUCTION + query for query in queries]
 
+print("Downloading images...")
+passages = [
+    Image.open(BytesIO(requests.get(
+        'https://github.com/OpenBMB/VisRAG/raw/refs/heads/master/scripts/demo/retriever/test_image/cat.jpeg'
+    ).content)).convert('RGB'),
+    Image.open(BytesIO(requests.get(
+        'https://github.com/OpenBMB/VisRAG/raw/refs/heads/master/scripts/demo/retriever/test_image/dog.jpg'
+    ).content)).convert('RGB')
+]
+print("Images downloaded.")
+
 embeddings_query = encode(queries)
 embeddings_doc = encode(passages)