Update app.py

app.py

@@ -1,15 +1,23 @@
+import os
 import streamlit as st
 import torch
 from transformers import BitsAndBytesConfig
 
-# Import llama-index and langchain modules
+# Import necessary modules from llama-index and langchain
 from llama_index.core import VectorStoreIndex, SimpleDirectoryReader, Settings, PromptTemplate
 from llama_index.llms.huggingface import HuggingFaceLLM
-from llama_index.llms.huggingface_api import HuggingFaceInferenceAPI
 from llama_index.embeddings.huggingface import HuggingFaceEmbedding
 from langchain.embeddings import HuggingFaceEmbeddings
 from llama_index.embeddings.langchain import LangchainEmbedding
 
+# ---------------------------
+# Retrieve Hugging Face Token from Environment Variables
+# ---------------------------
+hf_token = os.getenv("HF_TOKEN")
+if hf_token is None:
+    st.error("Missing Hugging Face token. Please set HF_TOKEN in your Space secrets.")
+    st.stop()
+
 # ---------------------------
 # Configure your LLM and embeddings
 # ---------------------------
@@ -25,7 +33,7 @@ quantization_config = BitsAndBytesConfig(
     bnb_4bit_compute_dtype=torch.float16
 )
 
-# Initialize the HuggingFaceLLM with your model settings
+# Initialize the HuggingFaceLLM with your model settings and authentication token
 llm = HuggingFaceLLM(
     context_window=4096,
     max_new_tokens=256,
@@ -37,7 +45,8 @@ llm = HuggingFaceLLM(
     device_map="auto",
     model_kwargs={
         "torch_dtype": torch.float16,
-        "quantization_config": quantization_config
+        "quantization_config": quantization_config,
+        "use_auth_token": hf_token  # Pass the HF token for gated access
     }
 )
 
@@ -53,8 +62,8 @@ Settings.chunk_size = 1024
 # ---------------------------
 # Load documents from repository
 # ---------------------------
-# The "data" folder should be part of your repository with your documents.
-DATA_DIR = "data"  # Ensure this folder exists and contains your documents.
+DATA_DIR = "data"  # Ensure this folder exists in your repository and contains your documents
+
 try:
     documents = SimpleDirectoryReader(DATA_DIR).load_data()
 except Exception as e:
@@ -63,20 +72,21 @@ except Exception as e:
 
 if not documents:
     st.warning("No documents found in the data folder. Please add your documents and redeploy.")
+    st.stop()
 else:
-    # Create the vector store index
+    # Create the vector store index and query engine
     index = VectorStoreIndex.from_documents(documents)
     query_engine = index.as_query_engine()
 
-    # ---------------------------
-    # Streamlit Interface
-    # ---------------------------
-    st.title("LLama Index Q&A Assistant")
+# ---------------------------
+# Streamlit Interface
+# ---------------------------
+st.title("LLama Index Q&A Assistant")
 
-    user_query = st.text_input("Enter your question:")
+user_query = st.text_input("Enter your question:")
 
-    if user_query:
-        with st.spinner("Querying..."):
-            response = query_engine.query(user_query)
-        st.markdown("### Response:")
-        st.write(response)
+if user_query:
+    with st.spinner("Querying..."):
+        response = query_engine.query(user_query)
+    st.markdown("### Response:")
+    st.write(response)