Update README.md

README.md

@@ -91,13 +91,13 @@ from models import CLIPVisionTower
 DEVICE = "cuda:0"
 PROMPT = "This is a dialog with AI assistant.\n"
 
-tokenizer = AutoTokenizer.from_pretrained("AIRI-Institute/OmniFusion", subfolder="OmniMistral-tokenizer", use_fast=False)
-model = AutoModelForCausalLM.from_pretrained("AIRI-Institute/OmniFusion", subfolder="OmniMistral-model", torch_dtype=torch.bfloat16, device_map=DEVICE)
+tokenizer = AutoTokenizer.from_pretrained("AIRI-Institute/OmniFusion", subfolder="OmniMistral-v1_1/tokenizer", use_fast=False)
+model = AutoModelForCausalLM.from_pretrained("AIRI-Institute/OmniFusion", subfolder="OmniMistral-v1_1/tuned-model", torch_dtype=torch.bfloat16, device_map=DEVICE)
 
-hf_hub_download(repo_id="AIRI-Institute/OmniFusion", filename="projection", local_dir='./')
-hf_hub_download(repo_id="AIRI-Institute/OmniFusion", filename="special_embeddings.pt", local_dir='./')
-projection = torch.load("projection", map_location=DEVICE)
-special_embs = torch.load("special_embeddings.pt", map_location=DEVICE)
+hf_hub_download(repo_id="AIRI-Institute/OmniFusion", filename="OmniMistral-v1_1/projection.pt", local_dir='./')
+hf_hub_download(repo_id="AIRI-Institute/OmniFusion", filename="OmniMistral-v1_1/special_embeddings.pt", local_dir='./')
+projection = torch.load("OmniMistral-v1_1/projection.pt", map_location=DEVICE)
+special_embs = torch.load("OmniMistral-v1_1/special_embeddings.pt", map_location=DEVICE)
 
 clip = CLIPVisionTower("openai/clip-vit-large-patch14-336")
 clip.load_model()
@@ -149,7 +149,7 @@ def gen_answer(model, tokenizer, clip, projection, query, special_embs, image=No
     return generated_texts
 
 img_url = "https://i.pinimg.com/originals/32/c7/81/32c78115cb47fd4825e6907a83b7afff.jpg"
-question = "who is the author?"
+question = "What is the sky color on this image?"
 img = Image.open(urlopen(img_url))
 
 answer = gen_answer(