Update README.md

README.md

@@ -82,9 +82,13 @@ def prepare_inputs(text, has_image=False, device='cuda'):
         tokenize=False
     )
 
-    text_chunks = [tokenizer(chunk).input_ids for chunk in inputs_formatted.split('<|image|>')]
-    input_ids = torch.tensor(text_chunks[0] + [-200] + text_chunks[1][1:], dtype=torch.long).unsqueeze(0).to(device)
-    attention_mask = torch.ones_like(input_ids).to(device)
+    if has_image:
+        text_chunks = [tokenizer(chunk).input_ids for chunk in inputs_formatted.split('<|image|>')]
+        input_ids = torch.tensor(text_chunks[0] + [-200] + text_chunks[1][1:], dtype=torch.long).unsqueeze(0).to(device)
+        attention_mask = torch.ones_like(input_ids).to(device)
+    else:
+        input_ids = torch.tensor(tokenizer(inputs_formatted).input_ids, dtype=torch.long).unsqueeze(0).to(device)
+        attention_mask = torch.ones_like(input_ids).to(device)
 
     return input_ids, attention_mask