Update app.py

app.py

@@ -3,29 +3,25 @@ import os
 import requests
 from langdetect import detect
 from PyPDF2 import PdfReader
-
-# Load the Hugging Face token from environment variables (secrets)
-token = os.environ.get("Key2")  # Replace "KEY2" with your secret key name
-
-# Function to query the Hugging Face API
-def query_huggingface_api(prompt, max_new_tokens=1000, temperature=0.7, top_k=50):
-    model_name = "HuggingFaceH4/zephyr-7b-alpha"  # Replace with your preferred model
-    api_url = f"https://api-inference.huggingface.co/models/{model_name}"
-    headers = {"Authorization": f"Bearer {token}"}
-    payload = {
-        "inputs": prompt,
-        "parameters": {
-            "max_new_tokens": max_new_tokens,
-            "temperature": temperature,
-            "top_k": top_k,
-        },
+import replicate  # For interacting with Llama models hosted on Replicate
+
+# Load the Replicate API token from environment variables
+replicate_api_token = os.environ.get("Key2")  # Replace with your Replicate API token
+
+# Function to query the Llama 3.2 7B Instruct model via Replicate
+def query_llama_model(prompt, max_new_tokens=1000, temperature=0.7, top_k=50):
+    model_name = "meta/llama-3-7b-instruct"  # Replace with the correct model name on Replicate
+    input_data = {
+        "prompt": prompt,
+        "max_new_tokens": max_new_tokens,
+        "temperature": temperature,
+        "top_k": top_k,
     }
-    response = requests.post(api_url, headers=headers, json=payload)
-    if response.status_code == 200:
-        return response.json()[0]["generated_text"]
-    else:
-        st.error(f"Error: {response.status_code} - {response.text}")
-        return None
+    response = replicate.run(
+        model_name,
+        input=input_data
+    )
+    return "".join(response)  # Replicate returns a generator, so we join it into a single string
 
 # Function to detect language
 def detect_language(text):
@@ -48,6 +44,20 @@ def extract_text_from_pdf(pdf_file):
             })
     return text_data
 
+# Function to search for query in PDF content
+def search_pdf_content(pdf_text_data, query):
+    results = []
+    for entry in pdf_text_data:
+        if query.lower() in entry["content"].lower():
+            results.append(entry)
+    return results
+
+# Function to split text into chunks
+def split_text_into_chunks(text, chunk_size=500):
+    words = text.split()
+    chunks = [" ".join(words[i:i + chunk_size]) for i in range(0, len(words), chunk_size)]
+    return chunks
+
 # Default system prompts for each query translation method
 DEFAULT_SYSTEM_PROMPTS = {
     "Multi-Query": """You are an AI language model assistant. Your task is to generate five 
@@ -120,33 +130,65 @@ def main():
             formatted_prompt = system_prompt.format(question=prompt)
             st.write("**Formatted System Prompt:**", formatted_prompt)
 
-            # Query the Hugging Face API for query translation
-            translated_queries = query_huggingface_api(formatted_prompt, max_new_tokens, temperature, top_k)
+            # Query the Llama model for query translation
+            translated_queries = query_llama_model(formatted_prompt, max_new_tokens, temperature, top_k)
             if translated_queries:
-                st.write("**Translated Queries:**", translated_queries)
+                st.write("**Translated Queries:**")
+                st.write(translated_queries.split("\n")[-1])  # Print only the updated question part
 
         # Indexing
         if st.button("Apply Indexing"):
             st.write(f"**Applied Indexing Method:** {indexing_method}")
-            # Implement indexing logic here
-            # Example: Indexing with ColBERT
-            if indexing_method == "ColBERT":
-                st.write("Indexing with ColBERT...")
-
-        # Query the Hugging Face API for final response
+            if pdf_file is not None:
+                # Extract and search PDF content
+                pdf_text_data = extract_text_from_pdf(pdf_file)
+                search_results = search_pdf_content(pdf_text_data, prompt)
+
+                if search_results:
+                    st.write("**Relevant Content from PDF:**")
+                    for result in search_results:
+                        st.write(f"**Page {result['page']}, Line {result['line']}:** {result['content']}")
+
+                    # Split text into chunks
+                    chunks = split_text_into_chunks("\n".join([result["content"] for result in search_results]))
+                    st.write("**Chunks Obtained from PDF:**")
+                    for i, chunk in enumerate(chunks):
+                        st.write(f"**Chunk {i + 1}:** {chunk}")
+
+                    # Print summary of split for Multi-Representation
+                    if indexing_method == "Multi-Representation":
+                        st.write("**Summary of Split:**")
+                        summary = query_llama_model(f"Summarize the following text:\n{chunks[0]}", max_new_tokens, temperature, top_k)
+                        st.write(summary)
+                else:
+                    st.write("**No relevant content found in the PDF.**")
+            else:
+                st.write("**No PDF uploaded.**")
+
+        # Generate Response
         if st.button("Generate Response"):
-            response = query_huggingface_api(prompt, max_new_tokens, temperature, top_k)
+            if pdf_file is not None:
+                # Extract and search PDF content
+                pdf_text_data = extract_text_from_pdf(pdf_file)
+                search_results = search_pdf_content(pdf_text_data, prompt)
+
+                if search_results:
+                    st.write("**Relevant Content from PDF:**")
+                    for result in search_results:
+                        st.write(f"**Page {result['page']}, Line {result['line']}:** \"{result['content']}\"")
+
+                    # Generate response based on PDF content
+                    pdf_context = "\n".join([result["content"] for result in search_results])
+                    response = query_llama_model(f"Based on the following context:\n{pdf_context}\n\nAnswer this question: {prompt}", max_new_tokens, temperature, top_k)
+                else:
+                    st.write("**No relevant content found in the PDF. Generating response without PDF context.**")
+                    response = query_llama_model(prompt, max_new_tokens, temperature, top_k)
+            else:
+                st.write("**No PDF uploaded. Generating response without PDF context.**")
+                response = query_llama_model(prompt, max_new_tokens, temperature, top_k)
+
             if response:
                 st.write("**Response:**", response)
 
-    # Process PDF content if uploaded
-    if pdf_file is not None:
-        pdf_text_data = extract_text_from_pdf(pdf_file)
-        if prompt:
-            # Search for relevant content in the PDF
-            for entry in pdf_text_data:
-                if prompt.lower() in entry["content"].lower():
-                    st.write(f"**Page {entry['page']}, Line {entry['line']}:** {entry['content']}")
-
 if __name__ == "__main__":
     main()