Update app.py

app.py

@@ -1,28 +1,20 @@
 import streamlit as st
+import os
+from huggingface_hub import InferenceApi
 from PyPDF2 import PdfReader
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain_community.embeddings import HuggingFaceEmbeddings
 from langchain.vectorstores import FAISS
-from langchain.chains import RetrievalQA
-from langchain.llms import HuggingFacePipeline
-import torch
-from transformers import pipeline
 from langdetect import detect
 
-# Load a smaller LLM with customizable parameters
-def load_llm(temperature, top_k, max_new_tokens):
+# Load the Hugging Face token from environment variables (secrets)
+token = os.environ.get("KEY2")  # Replace "KEY2" with your secret key name
+
+# Initialize the Hugging Face Inference API
+def load_llm():
     model_name = "HuggingFaceH4/zephyr-7b-alpha"  # Replace with your preferred model
-    pipe = pipeline(
-        "text-generation",
-        model=model_name,
-        torch_dtype=torch.float16,
-        device_map="auto",
-        temperature=temperature,
-        top_k=top_k,
-        max_new_tokens=max_new_tokens,  # Use max_new_tokens instead of max_length
-    )
-    llm = HuggingFacePipeline(pipeline=pipe)
-    return llm
+    api = InferenceApi(repo_id=model_name, token=token)
+    return api
 
 # Extract text from PDF
 def extract_text_from_pdf(file):
@@ -51,25 +43,25 @@ def create_vector_store(chunks, indexing_method="multi-representation", **kwargs
         vector_store = FAISS.from_texts(chunks, embeddings)
     return vector_store
 
-# Query the PDF
-def query_pdf(vector_store, query, llm, query_method="multi-query", **kwargs):
-    if query_method == "multi-query":
-        # Implement Multi-Query logic here
-        qa = RetrievalQA.from_chain_type(llm=llm, chain_type="stuff", retriever=vector_store.as_retriever())
-    elif query_method == "rag-fusion":
-        # Implement RAG Fusion logic here
-        qa = RetrievalQA.from_chain_type(llm=llm, chain_type="stuff", retriever=vector_store.as_retriever())
-    elif query_method == "decomposition":
-        # Implement Decomposition logic here
-        qa = RetrievalQA.from_chain_type(llm=llm, chain_type="stuff", retriever=vector_store.as_retriever())
-    elif query_method == "step-back":
-        # Implement Step Back logic here
-        qa = RetrievalQA.from_chain_type(llm=llm, chain_type="stuff", retriever=vector_store.as_retriever())
-    elif query_method == "hyde":
-        # Implement HyDE logic here
-        qa = RetrievalQA.from_chain_type(llm=llm, chain_type="stuff", retriever=vector_store.as_retriever())
-    result = qa.run(query)
-    return result
+# Query the PDF using the Hugging Face API
+def query_pdf(vector_store, query, api, query_method="multi-query", max_new_tokens=200, temperature=0.7, top_k=50):
+    # Retrieve relevant chunks from the vector store
+    docs = vector_store.similarity_search(query)
+    context = " ".join([doc.page_content for doc in docs])
+
+    # Create a prompt for the LLM
+    prompt = f"Context: {context}\n\nQuestion: {query}\n\nAnswer:"
+
+    # Query the Hugging Face API
+    response = api(
+        inputs=prompt,
+        parameters={
+            "max_new_tokens": max_new_tokens,
+            "temperature": temperature,
+            "top_k": top_k,
+        },
+    )
+    return response[0]["generated_text"], docs
 
 # Detect language of the text
 def detect_language(text):
@@ -115,9 +107,9 @@ def main():
             st.session_state.vector_store = create_vector_store(st.session_state.chunks, indexing_method=indexing_method)
             st.success("Vector store created!")
 
-    # Step 3: Load LLM with user-defined parameters
-    if "llm" not in st.session_state:
-        st.session_state.llm = None
+    # Step 3: Load LLM (Hugging Face API)
+    if "api" not in st.session_state:
+        st.session_state.api = None
 
     if st.session_state.vector_store:
         st.subheader("LLM Parameters")
@@ -125,11 +117,12 @@ def main():
         top_k = st.slider("Top-k", 1, 100, 50, help="Limits sampling to the top-k tokens.")
         max_new_tokens = st.slider("Max New Tokens", 50, 500, 200, help="Maximum number of tokens to generate.")
         if st.button("Load LLM"):
-            st.session_state.llm = load_llm(temperature=temperature, top_k=top_k, max_new_tokens=max_new_tokens)
+            api = load_llm()
+            st.session_state.api = api
             st.success("LLM loaded!")
 
     # Step 4: Query the PDF
-    if st.session_state.llm:
+    if st.session_state.api:
         st.subheader("Query Translation Options")
         query_method = st.selectbox(
             "Query Translation Method",
@@ -138,9 +131,19 @@ def main():
         )
         query = st.text_input("Ask a question about the PDF:")
         if query:
-            result = query_pdf(st.session_state.vector_store, query, st.session_state.llm, query_method=query_method)
-            st.write("**Answer:**", result["answer"])
-            st.write("**Source Text:**", result["source_text"])
+            answer, source_docs = query_pdf(
+                st.session_state.vector_store,
+                query,
+                st.session_state.api,
+                query_method=query_method,
+                max_new_tokens=max_new_tokens,
+                temperature=temperature,
+                top_k=top_k,
+            )
+            st.write("**Answer:**", answer)
+            st.write("**Source Text:**")
+            for doc in source_docs:
+                st.write(doc.page_content)
 
 if __name__ == "__main__":
     main()