Update README.md

README.md

@@ -42,12 +42,6 @@ from llava.mm_utils import tokenizer_image_token, process_images
 model_path = "shinyice/chatvector-llava-v1.5-plus-houou-v3-7b"
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-image_url = "https://huggingface.co/rinna/bilingual-gpt-neox-4b-minigpt4/resolve/main/sample.jpg"
-
-temperature = 0.0
-top_p = 1.0
-max_new_tokens = 256
-
 model = LlavaLlamaForCausalLM.from_pretrained(
     model_path,
     device_map=device,
@@ -69,6 +63,7 @@ eos_token_id_list = [
     tokenizer.bos_token_id,
 ]
 
+image_url = "https://huggingface.co/rinna/bilingual-gpt-neox-4b-minigpt4/resolve/main/sample.jpg"
 image = Image.open(requests.get(image_url, stream=True).raw).convert('RGB')
 
 if not isinstance(image, list):
@@ -84,7 +79,8 @@ else:
 
 image_sizes_tensor = torch.tensor(image_sizes, dtype=torch.int32, device=device)
 
-conv_mode = "v1"
+
+conv_mode = "v1" #mistral_instruct
 conv = conv_templates[conv_mode].copy()
 prompt = "猫の隣には何がありますか？"
 inp = DEFAULT_IMAGE_TOKEN + '\n' + prompt
@@ -98,6 +94,12 @@ input_ids = tokenizer_image_token(
     IMAGE_TOKEN_INDEX,
     return_tensors='pt'
 ).unsqueeze(0)
+if device == "cuda":
+    input_ids = input_ids.to(device)
+
+temperature = 0.0
+top_p = 1.0
+max_new_tokens = 256
 
 with torch.inference_mode():
     output = model.generate(