Update app.py

app.py

@@ -114,22 +114,22 @@ def generate_answers(image_path, question):
 @spaces.GPU
 def generate_caption_llava(image_path, caption_bool):
     if caption_bool:
-        text_prompt =f"[INST]  <image>\nCaption this image in detail in objective manner.[/INST]"
+        text_prompt ="[INST]  <image>\nCaption this image in detail in objective manner.[/INST]"
     else:
-        text_prompt =f"[INST]  <image>\nCaption this image briefly in objective manner. [/INST]"
+        text_prompt ="[INST]  <image>\nCaption this image briefly in objective manner. [/INST]"
     
     inputs = processor(text_prompt, Image.open(image_path), return_tensors="pt").to(device)
-
     output = model.generate(**inputs, max_new_tokens=100)
-
-    return processor.decode(output[0], skip_special_tokens=True)[len(text_prompt)-1:]
+    
+    return processor.decode(output[0], skip_special_tokens=True)).split("[/INST]")[1] 
 
 @spaces.GPU
 def generate_answer_llava(image_path, question):
     text_prompt =f"[INST]  <image>\n{question} [/INST]"
     inputs = processor(text_prompt, Image.open(image_path), return_tensors="pt").to(device)
     output = model.generate(**inputs, max_new_tokens=100)
-    return processor.decode(output[0], skip_special_tokens=True)[len(text_prompt)-1:]
+    
+    return processor.decode(output[0], skip_special_tokens=True).split("[/INST]")[1]