Update app.py

app.py

@@ -10,7 +10,7 @@ from transformers import pipeline
 from langdetect import detect
 
 # Load a smaller LLM with customizable parameters
-def load_llm(temperature, top_k, max_length):
+def load_llm(temperature, top_k, max_new_tokens):
     model_name = "HuggingFaceH4/zephyr-7b-alpha"  # Replace with your preferred model
     pipe = pipeline(
         "text-generation",
@@ -19,7 +19,7 @@ def load_llm(temperature, top_k, max_length):
         device_map="auto",
         temperature=temperature,
         top_k=top_k,
-        max_length=max_length,
+        max_new_tokens=max_new_tokens,  # Use max_new_tokens instead of max_length
     )
     llm = HuggingFacePipeline(pipeline=pipe)
     return llm
@@ -39,7 +39,7 @@ def split_text(text, chunk_size=1000, chunk_overlap=200):
     return chunks
 
 # Create embeddings and vector store
-def create_vector_store(chunks, indexing_method="multi-representation"):
+def create_vector_store(chunks, indexing_method="multi-representation", **kwargs):
     embeddings = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")
     if indexing_method == "multi-representation":
         vector_store = FAISS.from_texts(chunks, embeddings)
@@ -52,7 +52,7 @@ def create_vector_store(chunks, indexing_method="multi-representation"):
     return vector_store
 
 # Query the PDF
-def query_pdf(vector_store, query, llm, query_method="multi-query"):
+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())
@@ -89,53 +89,58 @@ def main():
         st.info("Using default PDF.")
         uploaded_file = "default.pdf"  # Add a default PDF
 
-    # Extract text
-    text = extract_text_from_pdf(uploaded_file)
-
-    # Detect language
-    language = detect_language(text)
-    st.write(f"Detected Language: {language}")
-
-    # Split text into chunks
-    chunk_size = st.slider("Chunk Size", 500, 2000, 1000)
-    chunk_overlap = st.slider("Chunk Overlap", 0, 500, 200)
-    chunks = split_text(text, chunk_size=chunk_size, chunk_overlap=chunk_overlap)
-
-    # Indexing options
-    indexing_method = st.selectbox(
-        "Indexing Method",
-        ["multi-representation", "raptors", "colbert"],
-        help="Choose how to index the PDF text."
-    )
-    st.write(f"**Indexing Method:** {indexing_method}")
-
-    # Create vector store
-    vector_store = create_vector_store(chunks, indexing_method=indexing_method)
-
-    # LLM Parameters
-    st.sidebar.header("LLM Parameters")
-    temperature = st.sidebar.slider("Temperature", 0.1, 1.0, 0.7, help="Controls randomness in the output.")
-    top_k = st.sidebar.slider("Top-k", 1, 100, 50, help="Limits sampling to the top-k tokens.")
-    max_length = st.sidebar.slider("Max Length", 50, 500, 200, help="Maximum length of the generated response.")
-
-    # Load LLM with user-defined parameters
-    llm = load_llm(temperature=temperature, top_k=top_k, max_length=max_length)
-
-    # Query translation options
-    query_method = st.selectbox(
-        "Query Translation Method",
-        ["multi-query", "rag-fusion", "decomposition", "step-back", "hyde"],
-        help="Choose a method to improve query retrieval."
-    )
-    st.write(f"**Query Translation Method:** {query_method}")
-
-    # User input
-    query = st.text_input("Ask a question about the PDF:")
-    if query:
-        # Query the PDF
-        result = query_pdf(vector_store, query, llm, query_method=query_method)
-        st.write("**Answer:**", result["answer"])
-        st.write("**Source Text:**", result["source_text"])
+    # Step 1: Extract text and split into chunks
+    if "text" not in st.session_state:
+        st.session_state.text = None
+    if "chunks" not in st.session_state:
+        st.session_state.chunks = None
+
+    if st.button("Extract Text and Split into Chunks"):
+        st.session_state.text = extract_text_from_pdf(uploaded_file)
+        st.session_state.chunks = split_text(st.session_state.text)
+        st.success("Text extracted and split into chunks!")
+
+    # Step 2: Create vector store
+    if "vector_store" not in st.session_state:
+        st.session_state.vector_store = None
+
+    if st.session_state.chunks:
+        st.subheader("Indexing Options")
+        indexing_method = st.selectbox(
+            "Indexing Method",
+            ["multi-representation", "raptors", "colbert"],
+            help="Choose how to index the PDF text."
+        )
+        if st.button("Create Vector Store"):
+            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
+
+    if st.session_state.vector_store:
+        st.subheader("LLM Parameters")
+        temperature = st.slider("Temperature", 0.1, 1.0, 0.7, help="Controls randomness in the output.")
+        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)
+            st.success("LLM loaded!")
+
+    # Step 4: Query the PDF
+    if st.session_state.llm:
+        st.subheader("Query Translation Options")
+        query_method = st.selectbox(
+            "Query Translation Method",
+            ["multi-query", "rag-fusion", "decomposition", "step-back", "hyde"],
+            help="Choose a method to improve query retrieval."
+        )
+        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"])
 
 if __name__ == "__main__":
     main()