Update app.py

app.py

@@ -12,11 +12,19 @@ from langchain_community.document_loaders import PyPDFLoader
 from langchain_community.embeddings import HuggingFaceEmbeddings
 from llama_parse import LlamaParse
 from langchain_core.documents import Document
+from huggingface_hub import InferenceClient
+import inspect
 
 # Environment variables and configurations
 huggingface_token = os.environ.get("HUGGINGFACE_TOKEN")
 llama_cloud_api_key = os.environ.get("LLAMA_CLOUD_API_KEY")
 
+MODELS = [
+    "mistralai/Mistral-7B-Instruct-v0.3",
+    "mistralai/Mixtral-8x7B-Instruct-v0.1",
+    "microsoft/Phi-3-mini-4k-instruct"
+]
+
 # Initialize LlamaParse
 llama_parser = LlamaParse(
     api_key=llama_cloud_api_key,
@@ -26,7 +34,7 @@ llama_parser = LlamaParse(
     language="en",
 )
 
-def load_document(file: NamedTemporaryFile, parser: str = "pypdf") -> List[Document]:
+def load_document(file: NamedTemporaryFile, parser: str = "llamaparse") -> List[Document]:
     """Loads and splits the document into pages."""
     if parser == "pypdf":
         loader = PyPDFLoader(file.name)
@@ -69,53 +77,55 @@ def update_vectors(files, parser):
     
     return f"Vector store updated successfully. Processed {total_chunks} chunks from {len(files)} files using {parser}."
 
-def generate_chunked_response(prompt, max_tokens=1000, max_chunks=5, temperature=0.7, repetition_penalty=1.1):
-    API_URL = "https://api-inference.huggingface.co/models/mistralai/Mistral-7B-Instruct-v0.3"
-    headers = {"Authorization": f"Bearer {huggingface_token}"}
-    payload = {
-        "inputs": prompt,
-        "parameters": {
-            "max_new_tokens": max_tokens,
-            "temperature": temperature,
-            "top_p": 0.4,
-            "top_k": 40,
-            "repetition_penalty": repetition_penalty,
-            "stop": ["</s>", "[/INST]"]
-        }
-    }
-    
+def generate_chunked_response(prompt, model, max_tokens=1000, num_calls=3, temperature=0.2, should_stop=False):
+    print(f"Starting generate_chunked_response with {num_calls} calls")
+    client = InferenceClient(model, token=huggingface_token)
     full_response = ""
-    for _ in range(max_chunks):
-        response = requests.post(API_URL, headers=headers, json=payload)
-        if response.status_code == 200:
-            result = response.json()
-            if isinstance(result, list) and len(result) > 0:
-                chunk = result[0].get('generated_text', '')
-                
-                # Remove any part of the chunk that's already in full_response
-                new_content = chunk[len(full_response):].strip()
-                
-                if not new_content:
-                    break  # No new content, so we're done
-                
-                full_response += new_content
-                
-                if chunk.endswith((".", "!", "?", "</s>", "[/INST]")):
-                    break
-                
-                # Update the prompt for the next iteration
-                payload["inputs"] = full_response
-            else:
-                break
-        else:
+    messages = [{"role": "user", "content": prompt}]
+    
+    for i in range(num_calls):
+        print(f"Starting API call {i+1}")
+        if should_stop:
+            print("Stop clicked, breaking loop")
             break
+        try:
+            for message in client.chat_completion(
+                messages=messages,
+                max_tokens=max_tokens,
+                temperature=temperature,
+                stream=True,
+            ):
+                if should_stop:
+                    print("Stop clicked during streaming, breaking")
+                    break
+                if message.choices and message.choices[0].delta and message.choices[0].delta.content:
+                    chunk = message.choices[0].delta.content
+                    full_response += chunk
+            print(f"API call {i+1} completed")
+        except Exception as e:
+            print(f"Error in generating response: {str(e)}")
     
     # Clean up the response
     clean_response = re.sub(r'<s>\[INST\].*?\[/INST\]\s*', '', full_response, flags=re.DOTALL)
     clean_response = clean_response.replace("Using the following context:", "").strip()
     clean_response = clean_response.replace("Using the following context from the PDF documents:", "").strip()
     
-    return clean_response
+    # Remove duplicate paragraphs and sentences
+    paragraphs = clean_response.split('\n\n')
+    unique_paragraphs = []
+    for paragraph in paragraphs:
+        if paragraph not in unique_paragraphs:
+            sentences = paragraph.split('. ')
+            unique_sentences = []
+            for sentence in sentences:
+                if sentence not in unique_sentences:
+                    unique_sentences.append(sentence)
+            unique_paragraphs.append('. '.join(unique_sentences))
+    
+    final_response = '\n\n'.join(unique_paragraphs)
+    
+    print(f"Final clean response: {final_response[:100]}...")
+    return final_response
 
 def duckduckgo_search(query):
     with DDGS() as ddgs:
@@ -127,113 +137,171 @@ class CitingSources(BaseModel):
         ...,
         description="List of sources to cite. Should be an URL of the source."
     )
+def chatbot_interface(message, history, use_web_search, model, temperature, num_calls):
+    if not message.strip():
+        return "", history
 
-def get_response_from_pdf(query, temperature=0.7, repetition_penalty=1.1):
-    embed = get_embeddings()
-    if os.path.exists("faiss_database"):
-        database = FAISS.load_local("faiss_database", embed, allow_dangerous_deserialization=True)
-    else:
-        return "No documents available. Please upload PDF documents to answer questions."
-
-    retriever = database.as_retriever()
-    relevant_docs = retriever.get_relevant_documents(query)
-    context_str = "\n".join([doc.page_content for doc in relevant_docs])
-
-    prompt = f"""<s>[INST] Using the following context from the PDF documents:
-{context_str}
-Write a detailed and complete response that answers the following user question: '{query}'
-Do not include a list of sources in your response. [/INST]"""
+    history = history + [(message, "")]
 
-    generated_text = generate_chunked_response(prompt, temperature=temperature, repetition_penalty=repetition_penalty)
+    try:
+        if use_web_search:
+            for main_content, sources in get_response_with_search(message, model, num_calls=num_calls, temperature=temperature):
+                history[-1] = (message, f"{main_content}\n\n{sources}")
+                yield history
+        else:
+            for partial_response in get_response_from_pdf(message, model, num_calls=num_calls, temperature=temperature):
+                history[-1] = (message, partial_response)
+                yield history
+    except gr.CancelledError:
+        yield history
 
-    # Clean the response
-    clean_text = re.sub(r'<s>\[INST\].*?\[/INST\]\s*', '', generated_text, flags=re.DOTALL)
-    clean_text = clean_text.replace("Using the following context from the PDF documents:", "").strip()
+def retry_last_response(history, use_web_search, model, temperature, num_calls):
+    if not history:
+        return history
+    
+    last_user_msg = history[-1][0]
+    history = history[:-1]  # Remove the last response
+    
+    return chatbot_interface(last_user_msg, history, use_web_search, model, temperature, num_calls)
 
-    return clean_text
+def respond(message, history, model, temperature, num_calls, use_web_search):
+    if use_web_search:
+        for main_content, sources in get_response_with_search(message, model, num_calls=num_calls, temperature=temperature):
+            yield f"{main_content}\n\n{sources}"
+    else:
+        for partial_response in get_response_from_pdf(message, model, num_calls=num_calls, temperature=temperature):
+            yield partial_response
 
-def get_response_with_search(query, temperature=0.7, repetition_penalty=1.1):
+def get_response_with_search(query, model, num_calls=3, temperature=0.2):
     search_results = duckduckgo_search(query)
     context = "\n".join(f"{result['title']}\n{result['body']}\nSource: {result['href']}\n" 
                         for result in search_results if 'body' in result)
     
-    prompt = f"""<s>[INST] Using the following context:
+    prompt = f"""Using the following context:
 {context}
 Write a detailed and complete research document that fulfills the following user request: '{query}'
-After writing the document, please provide a list of sources used in your response. [/INST]"""
+After writing the document, please provide a list of sources used in your response."""
     
-    generated_text = generate_chunked_response(prompt, temperature=temperature, repetition_penalty=repetition_penalty)
+    client = InferenceClient(model, token=huggingface_token)
     
-    # Clean the response
-    clean_text = re.sub(r'<s>\[INST\].*?\[/INST\]\s*', '', generated_text, flags=re.DOTALL)
-    clean_text = clean_text.replace("Using the following context:", "").strip()
-    
-    # Split the content and sources
-    parts = clean_text.split("Sources:", 1)
-    main_content = parts[0].strip()
-    sources = parts[1].strip() if len(parts) > 1 else ""
+    main_content = ""
+    for i in range(num_calls):
+        for message in client.chat_completion(
+            messages=[{"role": "user", "content": prompt}],
+            max_tokens=1000,
+            temperature=temperature,
+            stream=True,
+        ):
+            if message.choices and message.choices[0].delta and message.choices[0].delta.content:
+                chunk = message.choices[0].delta.content
+                main_content += chunk
+                yield main_content, ""  # Yield partial main content without sources
+
+def get_response_from_pdf(query, model, num_calls=3, temperature=0.2):
+    embed = get_embeddings()
+    if os.path.exists("faiss_database"):
+        database = FAISS.load_local("faiss_database", embed, allow_dangerous_deserialization=True)
+    else:
+        yield "No documents available. Please upload PDF documents to answer questions."
+        return
+
+    retriever = database.as_retriever()
+    relevant_docs = retriever.get_relevant_documents(query)
+    context_str = "\n".join([doc.page_content for doc in relevant_docs])
+
+    prompt = f"""Using the following context from the PDF documents:
+{context_str}
+Write a detailed and complete response that answers the following user question: '{query}'"""
+
+    client = InferenceClient(model, token=huggingface_token)
     
-    return main_content, sources
+    response = ""
+    for i in range(num_calls):
+        for message in client.chat_completion(
+            messages=[{"role": "user", "content": prompt}],
+            max_tokens=1000,
+            temperature=temperature,
+            stream=True,
+        ):
+            if message.choices and message.choices[0].delta and message.choices[0].delta.content:
+                chunk = message.choices[0].delta.content
+                response += chunk
+                yield response  # Yield partial response
 
-def chatbot_interface(message, history, use_web_search, temperature, repetition_penalty):
-    if use_web_search:
-        main_content, sources = get_response_with_search(message, temperature, repetition_penalty)
-        formatted_response = f"{main_content}\n\nSources:\n{sources}"
+def vote(data: gr.LikeData):
+    if data.liked:
+        print(f"You upvoted this response: {data.value}")
     else:
-        response = get_response_from_pdf(message, temperature, repetition_penalty)
-        formatted_response = response
+        print(f"You downvoted this response: {data.value}")
+
+css = """
+/* Add your custom CSS here */
+"""
 
-    history.append((message, formatted_response))
-    return history
+# Define the checkbox outside the demo block
+use_web_search = gr.Checkbox(label="Use Web Search", value=False)
 
-# Gradio interface
-with gr.Blocks() as demo:
-    gr.Markdown("# AI-powered Web Search and PDF Chat Assistant")
+demo = gr.ChatInterface(
+    respond,
+    additional_inputs=[
+        gr.Dropdown(choices=MODELS, label="Select Model", value=MODELS[0]),
+        gr.Slider(minimum=0.1, maximum=1.0, value=0.2, step=0.1, label="Temperature"),
+        gr.Slider(minimum=1, maximum=5, value=1, step=1, label="Number of API Calls"),
+        use_web_search  # Add this line to include the checkbox
+    ],
+    title="AI-powered Web Search and PDF Chat Assistant",
+    description="Chat with your PDFs or use web search to answer questions.",
+    theme=gr.themes.Soft(
+        primary_hue="orange",
+        secondary_hue="amber",
+        neutral_hue="gray",
+        font=[gr.themes.GoogleFont("Exo"), "ui-sans-serif", "system-ui", "sans-serif"]
+    ).set(
+        body_background_fill_dark="#0c0505",
+        block_background_fill_dark="#0c0505",
+        block_border_width="1px",
+        block_title_background_fill_dark="#1b0f0f",
+        input_background_fill_dark="#140b0b",
+        button_secondary_background_fill_dark="#140b0b",
+        border_color_accent_dark="#1b0f0f",
+        border_color_primary_dark="#1b0f0f",
+        background_fill_secondary_dark="#0c0505",
+        color_accent_soft_dark="transparent",
+        code_background_fill_dark="#140b0b"
+    ),
     
+    css=css,
+    examples=[
+        ["Tell me about the contents of the uploaded PDFs."],
+        ["What are the main topics discussed in the documents?"],
+        ["Can you summarize the key points from the PDFs?"]
+    ],
+    cache_examples=False,
+    analytics_enabled=False,
+)
+
+# Add file upload functionality
+with demo:
+    gr.Markdown("## Upload PDF Documents")
+
     with gr.Row():
         file_input = gr.Files(label="Upload your PDF documents", file_types=[".pdf"])
-        parser_dropdown = gr.Dropdown(choices=["pypdf", "llamaparse"], label="Select PDF Parser", value="pypdf")
+        parser_dropdown = gr.Dropdown(choices=["pypdf", "llamaparse"], label="Select PDF Parser", value="llamaparse")
         update_button = gr.Button("Upload Document")
     
     update_output = gr.Textbox(label="Update Status")
     update_button.click(update_vectors, inputs=[file_input, parser_dropdown], outputs=update_output)
-    
-    chatbot = gr.Chatbot(label="Conversation")
-    msg = gr.Textbox(label="Ask a question")
-    use_web_search = gr.Checkbox(label="Use Web Search", value=False)
-    
-    with gr.Row():
-        temperature_slider = gr.Slider(minimum=0.1, maximum=1.0, value=0.7, step=0.1, label="Temperature")
-        repetition_penalty_slider = gr.Slider(minimum=1.0, maximum=2.0, value=1.1, step=0.1, label="Repetition Penalty")
-    
-    submit = gr.Button("Submit")
-
-    gr.Examples(
-        examples=[
-            ["What are the latest developments in AI?"],
-            ["Tell me about recent updates on GitHub"],
-            ["What are the best hotels in Galapagos, Ecuador?"],
-            ["Summarize recent advancements in Python programming"],
-        ],
-        inputs=msg,
-    )
-
-    submit.click(chatbot_interface, 
-                 inputs=[msg, chatbot, use_web_search, temperature_slider, repetition_penalty_slider], 
-                 outputs=[chatbot])
-    msg.submit(chatbot_interface, 
-               inputs=[msg, chatbot, use_web_search, temperature_slider, repetition_penalty_slider], 
-               outputs=[chatbot])
 
+    
     gr.Markdown(
     """
     ## How to use
     1. Upload PDF documents using the file input at the top.
     2. Select the PDF parser (pypdf or llamaparse) and click "Upload Document" to update the vector store.
-    3. Ask questions in the textbox. 
+    3. Ask questions in the chat interface. 
     4. Toggle "Use Web Search" to switch between PDF chat and web search.
-    5. Adjust Temperature and Repetition Penalty sliders to fine-tune the response generation.
-    6. Click "Submit" or press Enter to get a response.
+    5. Adjust Temperature and Number of API Calls to fine-tune the response generation.
+    6. Use the provided examples or ask your own questions.
     """
     )