Update "Quick Start"

README.md

@@ -21,33 +21,47 @@ from transformers import (
     LlavaProcessor,
     LlavaForConditionalGeneration,
 )
-import torch
 from PIL import Image
 import requests
 
 MODEL_NAME = "Seungyoun/llava-llama-3-8b-hf"
 
 processor = LlavaProcessor.from_pretrained(MODEL_NAME)
-# add 128257 <image> , <pad>
-processor.tokenizer.add_tokens(["<|image|>", "<pad>"], special_tokens=True)
+processor.tokenizer.add_tokens(
+    ["<|image|>", "<pad>"], special_tokens=True
+)  # add 128257 <|image|> , <pad>
 
 model = LlavaForConditionalGeneration.from_pretrained(MODEL_NAME).to("cuda:0")
-# resize embeddings
-model.resize_token_embeddings(len(processor.tokenizer))
+model.resize_token_embeddings(
+    len(processor.tokenizer)
+)  # resize embeddings for new tokens
 
 
 # prepare image and text prompt, using the appropriate prompt template
-url = "https://encrypted-tbn0.gstatic.com/images?q=tbn:ANd9GcTd4g61TSw890IYKBbPMgXPyWAKdVOpWWUAF0-FGzgX2Q&s"
+url = "https://upload.wikimedia.org/wikipedia/commons/1/18/Kochendes_wasser02.jpg"
 image = Image.open(requests.get(url, stream=True).raw)
-prompt = "A chat between a curious human and an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the human's questions. USER: <|image|>\nWhat is shown in this image? ASSISTANT:" # FIX : Chat template
+
+template = """<|start_header_id|>system<|end_header_id|>{system_prompt}<|eot_id|>
+<|start_header_id|>user<|end_header_id|>{user_msg_1}<|eot_id|>
+<|start_header_id|>assistant<|end_header_id|>"""
+
+terminators = [
+    processor.tokenizer.eos_token_id,
+    processor.tokenizer.convert_tokens_to_ids("<|eot_id|>"),
+]
+
+prompt = template.format(
+    system_prompt="As a vision-llm, your task is to analyze and describe the contents of the image presented to you. Examine the photograph closely and provide a comprehensive, detailed caption. You should identify and describe the various food items and their arrangement, as well as any discernible textures, colors, and specific features of the containers they are in. Highlight the variety and how these contribute to the overall visual appeal of the meal. Your description should help someone who cannot see the image to visualize its contents accurately.",
+    user_msg_1="<|image|>\nGive me detailed description of the image.",
+)
 
 inputs = processor(prompt, image, return_tensors="pt").to("cuda:0")
 
 # autoregressively complete prompt
-output = model.generate(**inputs, max_new_tokens=100)
+output = model.generate(**inputs, max_new_tokens=1024, eos_token_id=terminators)
 
-print(processor.decode(output[0], skip_special_tokens=True))
-# What is shown in this image? ASSISTANT: The image shows a heartwarming scene of two dogs sitting together on a couch. The dogs are of different breeds, one being a golden retriever and the other being a tabby cat. The dogs are sitting close together, indicating a strong bond between them. The image captures a beautiful moment of companionship between two different species. sit on couch. golden retriever and tabby cat. dogs are sitting together. companionship between two different species.
+print(processor.decode(output[0], skip_special_tokens=False))
+# The image captures a moment in a kitchen. The main focus is a white electric kettle, which is plugged in and resting on a black stovetop. The stovetop has four burners, although only one is occupied by the kettle. The background is blurred, drawing attention to the kettle and stovetop. The image does not contain any text or additional objects. The relative position of the objects is such that the kettle is on the stovetop, and the background is blurred.
 ```
 ---