Update README.md

README.md

@@ -30,17 +30,19 @@ from PIL import Image
 
 import torch
 from transformers import AutoProcessor
-from heron.models.git_llm.git_llama import GitLlamaForCausalLM
+from heron.models.git_llm.git_japanese_stablelm_alpha import GitJapaneseStableLMAlphaForCausalLM
 
 device_id = 0
 
 # prepare a pretrained model
-model = GitLlamaForCausalLM.from_pretrained('turing-motors/heron-chat-git-ja-stablelm-base-7b-v0')
+model = GitJapaneseStableLMAlphaForCausalLM.from_pretrained(
+    'turing-motors/heron-chat-git-ja-stablelm-base-7b-v0', torch_dtype=torch.float16
+)
 model.eval()
 model.to(f"cuda:{device_id}")
 
 # prepare a processor
-processor = AutoProcessor.from_pretrained('turing-motors/heron-chat-git-ja-stablelm-base-7b-v0', additional_special_tokens=["▁▁"])
+processor = AutoProcessor.from_pretrained('turing-motors/heron-chat-git-ja-stablelm-base-7b-v0')
 
 # prepare inputs
 url = "https://www.barnorama.com/wp-content/uploads/2016/12/03-Confusing-Pictures.jpg"
@@ -68,7 +70,7 @@ with torch.no_grad():
     out = model.generate(**inputs, max_length=256, do_sample=False, temperature=0., eos_token_id=eos_token_id_list)
 
 # print result
-print(processor.tokenizer.batch_decode(out))
+print(processor.tokenizer.batch_decode(out)[0])
 ```