Update my_model/KBVQA.py

my_model/KBVQA.py

@@ -52,7 +52,10 @@ class KBVQA:
     def __init__(self):
 
        # self.col1, self.col2, self.col3 = st.columns([0.2, 0.6, 0.2])
-        self.kbvqa_model_name: str = config.KBVQA_MODEL_NAME
+        if st.session_state["method"] == "7b-Fine-Tuned Model":
+            self.kbvqa_model_name: str = config.KBVQA_MODEL_NAME_7b
+        elif st.session_state["method"] == "13b-Fine-Tuned Model":
+            self.kbvqa_model_name: str = config.KBVQA_MODEL_NAME_13b
         self.quantization: str = config.QUANTIZATION
         self.max_context_window: int = config.MAX_CONTEXT_WINDOW
         self.add_eos_token: bool = config.ADD_EOS_TOKEN