Update README.md

README.md

@@ -43,18 +43,22 @@ model = AutoModelForCausalLM.from_pretrained(
     trust_remote_code=True)
 tokenizer = AutoTokenizer.from_pretrained("ManishThota/Sparrow", trust_remote_code=True)
 
-#Set inputs
-text = "A chat between a curious user and an artificial intelligence assistant. USER: <image>\nCan you explain the slide? ASSISTANT:"
-image = Image.open("images/week_02_page_02")
-
-input_ids = tokenizer(text, return_tensors='pt').input_ids
-image_tensor = model.image_preprocess(image)
-
-#Generate the answer
-output_ids = model.generate(
-    input_ids,
-    max_new_tokens=1500,
-    images=image_tensor,
-    use_cache=True)[0]
-print(tokenizer.decode(output_ids[input_ids.shape[1]:], skip_special_tokens=True).strip())
+#function to generate the answer
+def predict(question, image_path):
+    #Set inputs
+    text = f"A chat between a curious user and an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the user's questions. USER: <image>\n{question}? ASSISTANT:"
+    image = Image.open(image_path)
+    
+    input_ids = tokenizer(text, return_tensors='pt').input_ids
+    image_tensor = model.image_preprocess(image)
+    
+    #Generate the answer
+    output_ids = model.generate(
+        input_ids,
+        max_new_tokens=25,
+        images=image_tensor,
+        use_cache=True)[0]
+    
+    return tokenizer.decode(output_ids[input_ids.shape[1]:], skip_special_tokens=True).strip()
+
 ```