Update my_model/KBVQA.py

my_model/KBVQA.py

@@ -224,7 +224,7 @@ class KBVQA:
         return p
 
     @staticmethod
-    def trim_objects(self, detected_objects_str):
+    def trim_objects(detected_objects_str):
         """
         Trim the last object from the detected objects string.
     
@@ -257,7 +257,9 @@ class KBVQA:
         prompt = self.format_prompt(question, caption=caption, objects=detected_objects_str)
         num_tokens = len(self.kbvqa_tokenizer.tokenize(prompt))
         self.current_prompt_length = num_tokens
-        
+        if self.current_prompt_length > self.max_context_window:
+            trim = True
+            st.warning(f"Prompt length is {self.current_prompt_length} which is larger than the maximum context window of LLaMA-2, objects detected with low confidence will be removed one at a time until the prompt length is within the maximum context window ...")
         while self.current_prompt_length > self.max_context_window:
             detected_objects_str = self.trim_objects(detected_objects_str)
             prompt = self.format_prompt(question, caption=caption, objects=detected_objects_str)
@@ -265,7 +267,9 @@ class KBVQA:
     
             if detected_objects_str == "":
                 break  # Break if no objects are left
-        
+        if trim:
+            st.warning(f"New prompt length is: {self.current_prompt_length}")
+            trim = False
 
         model_inputs = self.kbvqa_tokenizer(prompt, add_special_tokens=False, return_tensors="pt").to('cuda')
         free_gpu_resources()