Add HuggingFace authentication and better error handling for fine-tuned model

src/streamlit_app.py

@@ -45,11 +45,21 @@ load_dotenv()
 def load_fine_tuned_model():
     """Load the fine-tuned model from Hugging Face Hub"""
     try:
+        # Get HuggingFace token from environment
+        hf_token = os.getenv("HF_API_KEY")
+        if not hf_token:
+            st.error("HuggingFace API token not found. Please set HF_API_KEY in your environment.")
+            return None, None
+        
         # Replace with your actual repository name
         model_name = "kundan621/tinyllama-makemytrip-financial-qa"
         
-        # Load tokenizer
-        tokenizer = AutoTokenizer.from_pretrained(model_name)
+        # Load tokenizer with authentication
+        tokenizer = AutoTokenizer.from_pretrained(
+            model_name, 
+            token=hf_token,
+            trust_remote_code=True
+        )
         
         # Load base model
         base_model = AutoModelForCausalLM.from_pretrained(
@@ -59,12 +69,17 @@ def load_fine_tuned_model():
             trust_remote_code=True,
         )
         
-        # Load the fine-tuned PEFT model
-        model = PeftModel.from_pretrained(base_model, model_name)
+        # Load the fine-tuned PEFT model with authentication
+        model = PeftModel.from_pretrained(
+            base_model, 
+            model_name,
+            token=hf_token
+        )
         
         return model, tokenizer
     except Exception as e:
         st.error(f"Error loading fine-tuned model: {e}")
+        st.info("Make sure your model repository is public or you have the correct access permissions.")
         return None, None
 
 def generate_fine_tuned_response(model, tokenizer, question):
@@ -120,15 +135,38 @@ def generate_fine_tuned_response(model, tokenizer, question):
 st.set_page_config(page_title="Finance QA Assistant", layout="centered")
 st.title("Finance QA Assistant")
 
+# Add information about the modes
+with st.expander("ℹ️ About the Modes"):
+    st.markdown("""
+    **RAG Mode**: Uses Retrieval-Augmented Generation with a vector database and external LLM API.
+    
+    **Fine-Tuned Mode**: Uses a custom fine-tuned TinyLlama model (requires authentication).
+    
+    *Note: Fine-tuned mode requires a HuggingFace API token and access to the private model repository.*
+    """)
+
 # Load fine-tuned model if Fine-Tuned mode is available
 fine_tuned_model, fine_tuned_tokenizer = None, None
 
-mode = st.radio("Choose Answering Mode:", ["RAG", "Fine-Tuned"], horizontal=True)
+# Check if HuggingFace token is available
+hf_token = os.getenv("HF_API_KEY")
+if hf_token:
+    available_modes = ["RAG", "Fine-Tuned"]
+else:
+    available_modes = ["RAG"]
+    st.warning("⚠️ Fine-Tuned mode is not available. HuggingFace API token is required for accessing private models.")
+
+mode = st.radio("Choose Answering Mode:", available_modes, horizontal=True)
 
 if mode == "Fine-Tuned":
     if fine_tuned_model is None or fine_tuned_tokenizer is None:
         with st.spinner("Loading fine-tuned model..."):
             fine_tuned_model, fine_tuned_tokenizer = load_fine_tuned_model()
+            
+        # If model loading failed, fall back to RAG
+        if fine_tuned_model is None or fine_tuned_tokenizer is None:
+            st.error("Failed to load fine-tuned model. Falling back to RAG mode.")
+            mode = "RAG"
 
 query = st.text_input("Enter your question:")