Update my_model/KBVQA.py

my_model/KBVQA.py

@@ -34,6 +34,7 @@ class KBVQA:
         kbvqa_model (Optional[AutoModelForCausalLM]): The fine-tuned language model for KBVQA.
         bnb_config (BitsAndBytesConfig): Configuration for BitsAndBytes optimized model.
         access_token (str): Access token for Hugging Face API.
+        current_prompt_length (int): Prompt length.
 
     Methods:
         create_bnb_config: Creates a BitsAndBytes configuration based on the quantization setting.
@@ -66,6 +67,7 @@ class KBVQA:
         self.kbvqa_model: Optional[AutoModelForCausalLM] = None
         self.bnb_config: BitsAndBytesConfig = self.create_bnb_config()
         self.access_token: str = config.HUGGINGFACE_TOKEN
+        self.current_prompt_length = None
         
  
     def create_bnb_config(self) -> BitsAndBytesConfig:
@@ -227,8 +229,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 num_tokens > self.max_context_window:
-            st.write(f"Prompt too long with {num_tokens} tokens, consider increasing the confidence threshold for the object detector")
+            st.warning(f"Prompt too long with {num_tokens} tokens, consider increasing the confidence threshold for the object detector")
             return
 
         model_inputs = self.kbvqa_tokenizer(prompt, add_special_tokens=False, return_tensors="pt").to('cuda')