Update my_model/KBVQA.py

my_model/KBVQA.py

@@ -222,7 +222,22 @@ class KBVQA:
             p = f"""{history}\n{B_SENT}{B_INST} {B_QES}{current_query}{E_QES}{E_INST}"""
         
         return p
-       
+
+    @staticmethod
+    def trim_objects(self, detected_objects_str):
+        """
+        Trim the last object from the detected objects string.
+    
+        Args:
+        - detected_objects_str (str): String containing detected objects.
+    
+        Returns:
+        - (str): The string with the last object removed.
+        """
+        objects = detected_objects_str.strip().split("\n")
+        if len(objects) >= 1:
+            return "\n".join(objects[:-1])
+        return ""
 
     def generate_answer(self, question: str, caption: str, detected_objects_str: str) -> str:
         """
@@ -236,13 +251,21 @@ class KBVQA:
         Returns:
             str: The generated answer to the question.
         """
+    
+        
         free_gpu_resources()
         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 num_tokens > self.max_context_window:
-            st.warning(f"Prompt too long with {num_tokens} tokens, consider increasing the confidence threshold for the object detector")
-            return
+        
+        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)
+            self.current_prompt_length = len(self.kbvqa_tokenizer.tokenize(prompt))
+    
+            if detected_objects_str == "":
+                break  # Break if no objects are left
+        
 
         model_inputs = self.kbvqa_tokenizer(prompt, add_special_tokens=False, return_tensors="pt").to('cuda')
         free_gpu_resources()