Update my_model/KBVQA.py

my_model/KBVQA.py

@@ -48,24 +48,23 @@ class KBVQA():
     """
 
     def __init__(self):
-        
-        self.kbvqa_model_name = config.KBVQA_MODEL_NAME
-        self.quantization = config.QUANTIZATION
-        self.max_context_window = config.MAX_CONTEXT_WINDOW
-        self.add_eos_token = config.ADD_EOS_TOKEN
-        self.trust_remote = config.TRUST_REMOTE
-        self.use_fast = config.USE_FAST
-        self.low_cpu_mem_usage=config.LOW_CPU_MEM_USAGE
-        self.kbvqa_tokenizer = None
-        self.captioner = None
-        self.detector = None
-        self.detection_model = None
-        self.detection_confidence = None 
-        self.kbvqa_model = None
-        self.bnb_config = self.create_bnb_config()
-        self.access_token = config.HUGGINGFACE_TOKEN
  
-
+        self.model_name: str = config.KBVQA_MODEL_NAME
+        self.quantization: str = config.QUANTIZATION
+        self.max_context_window: int = config.MAX_CONTEXT_WINDOW
+        self.add_eos_token: bool = config.ADD_EOS_TOKEN
+        self.trust_remote: bool = config.TRUST_REMOTE
+        self.use_fast: bool = config.USE_FAST
+        self.low_cpu_mem_usage: bool = config.LOW_CPU_MEM_USAGE
+        self.kbvqa_tokenizer: Optional[AutoTokenizer] = None
+        self.captioner: Optional[ImageCaptioningModel] = None
+        self.detector: Optional[ObjectDetector] = None
+        self.detection_model: Optional[str] = None
+        self.detection_confidence: Optional[float] = None
+        self.kbvqa_model: Optional[AutoModelForCausalLM] = None
+        self.bnb_config: BitsAndBytesConfig = self.create_bnb_config()
+        self.access_token: str = config.HUGGINGFACE_TOKEN
+        
  
     def create_bnb_config(self) -> BitsAndBytesConfig:
         """