Update app.py

app.py

@@ -1,203 +1,103 @@
-import os
-import json
-import re
 import gradio as gr
+from huggingface_hub import InferenceApi
+from duckduckgo_search import DDGS
 import requests
-import random
-import urllib.parse
-from tempfile import NamedTemporaryFile
-from bs4 import BeautifulSoup
+import json
 from typing import List
 from pydantic import BaseModel, Field
-from huggingface_hub import InferenceApi
-from duckduckgo_search import DDGS
-from langchain_community.vectorstores import FAISS
-from langchain_community.document_loaders import PyPDFLoader
-from langchain_community.embeddings import HuggingFaceEmbeddings
-from langchain_community.llms import HuggingFaceHub
-from langchain_core.documents import Document
-from sentence_transformers import SentenceTransformer
-from llama_parse import LlamaParse
-import logging
-
-# Set up logging
-logging.basicConfig(level=logging.DEBUG, format='%(asctime)s - %(levelname)s - %(message)s')
+import os
 
 # Environment variables and configurations
 huggingface_token = os.environ.get("HUGGINGFACE_TOKEN")
-llama_cloud_api_key = os.environ.get("LLAMA_CLOUD_API_KEY")
-
-# Initialize SentenceTransformer and LlamaParse
-sentence_model = SentenceTransformer('paraphrase-MiniLM-L6-v2')
-llama_parser = LlamaParse(
-    api_key=llama_cloud_api_key,
-    result_type="markdown",
-    num_workers=4,
-    verbose=True,
-    language="en",
-)
-
-def load_document(file: NamedTemporaryFile, parser: str = "pypdf") -> List[Document]:
-    if parser == "pypdf":
-        loader = PyPDFLoader(file.name)
-        return loader.load_and_split()
-    elif parser == "llamaparse":
-        try:
-            documents = llama_parser.load_data(file.name)
-            return [Document(page_content=doc.text, metadata={"source": file.name}) for doc in documents]
-        except Exception as e:
-            print(f"Error using Llama Parse: {str(e)}")
-            print("Falling back to PyPDF parser")
-            loader = PyPDFLoader(file.name)
-            return loader.load_and_split()
-    else:
-        raise ValueError("Invalid parser specified. Use 'pypdf' or 'llamaparse'.")
-
-def update_vectors(files, parser):
-    if not files:
-        return "Please upload at least one PDF file."
-    
-    embed = HuggingFaceEmbeddings(model_name="sentence-transformers/all-mpnet-base-v2")
-    total_chunks = 0
-    
-    all_data = []
-    for file in files:
-        data = load_document(file, parser)
-        all_data.extend(data)
-        total_chunks += len(data)
-    
-    if os.path.exists("faiss_database"):
-        database = FAISS.load_local("faiss_database", embed, allow_dangerous_deserialization=True)
-        database.add_documents(all_data)
-    else:
-        database = FAISS.from_documents(all_data, embed)
-    
-    database.save_local("faiss_database")
-    
-    return f"Vector store updated successfully. Processed {total_chunks} chunks from {len(files)} files using {parser}."
-
-def clear_cache():
-    if os.path.exists("faiss_database"):
-        os.remove("faiss_database")
-        return "Cache cleared successfully."
-    else:
-        return "No cache to clear."
-
-def get_model(temperature, top_p, repetition_penalty):
-    return HuggingFaceHub(
-        repo_id="mistralai/Mistral-7B-Instruct-v0.3",
-        model_kwargs={
-            "temperature": temperature,
-            "top_p": top_p,
-            "repetition_penalty": repetition_penalty,
-            "max_length": 1000
-        },
-        huggingfacehub_api_token=huggingface_token
-    )
 
+# Function to perform a DuckDuckGo search
 def duckduckgo_search(query):
-    logging.debug(f"Performing DuckDuckGo search for query: {query}")
     with DDGS() as ddgs:
-        results = list(ddgs.text(query, max_results=5))
-    logging.debug(f"Search returned {len(results)} results")
+        results = ddgs.text(query, max_results=5)
     return results
 
-def get_response_with_search(query, temperature, top_p, repetition_penalty, use_pdf=False):
-    logging.debug(f"Getting response for query: {query}")
-    logging.debug(f"Parameters: temperature={temperature}, top_p={top_p}, repetition_penalty={repetition_penalty}, use_pdf={use_pdf}")
-    
-    model = get_model(temperature, top_p, repetition_penalty)
-    embed = HuggingFaceEmbeddings(model_name="sentence-transformers/all-mpnet-base-v2")
+class CitingSources(BaseModel):
+    sources: List[str] = Field(
+        ...,
+        description="List of sources to cite. Should be an URL of the source."
+    )
 
-    if use_pdf and os.path.exists("faiss_database"):
-        logging.debug("Using PDF database for context")
-        database = FAISS.load_local("faiss_database", embed, allow_dangerous_deserialization=True)
-        retriever = database.as_retriever()
-        relevant_docs = retriever.get_relevant_documents(query)
-        context = "\n".join([f"Content: {doc.page_content}\nSource: {doc.metadata['source']}\n" for doc in relevant_docs])
-    else:
-        logging.debug("Using web search for context")
-        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)
+def get_response_with_search(query):
+    # Perform the web search
+    search_results = duckduckgo_search(query)
     
-    logging.debug(f"Context generated. Length: {len(context)} characters")
-
+    # Use the search results as context for the model
+    context = "\n".join(f"{result['title']}\n{result['body']}\nSource: {result['href']}\n" 
+                        for result in search_results if 'body' in result)
+    
+    # Prompt formatted for Mistral-7B-Instruct
     prompt = f"""<s>[INST] Using the following context:
 {context}
 Write a detailed and complete research document that fulfills the following user request: '{query}'
-After the main content, provide a list of sources used in your response, prefixed with 'Sources:'. 
-Do not include any part of these instructions in your response. [/INST]"""
-
-    logging.debug("Sending prompt to model")
-    response = model(prompt)
-    logging.debug(f"Received response from model. Length: {len(response)} characters")
-    
-    main_content, sources = split_response(response)
-    
-    logging.debug(f"Split response. Main content length: {len(main_content)}, Sources length: {len(sources)}")
-    return main_content, sources
-
-def split_response(response):
-    logging.debug("Splitting response")
-    logging.debug(f"Original response: {response[:100]}...") # Log first 100 characters
-    
-    # Remove any remaining instruction text
-    response = re.sub(r'\[/?INST\]', '', response)
-    response = re.sub(r'~~.*?~~', '', response, flags=re.DOTALL)
-    
-    logging.debug(f"After removing instructions: {response[:100]}...") # Log first 100 characters
-    
-    # Split the response into main content and sources
-    parts = response.split("Sources:", 1)
-    main_content = parts[0].strip()
-    sources = parts[1].strip() if len(parts) > 1 else ""
-    
-    logging.debug(f"Main content starts with: {main_content[:100]}...") # Log first 100 characters
-    logging.debug(f"Sources: {sources[:100]}...") # Log first 100 characters
-    
-    return main_content, sources
+After writing the document, please provide a list of sources used in your response. [/INST]"""
+    
+    # API endpoint for Mistral-7B-Instruct-v0.3
+    API_URL = "https://api-inference.huggingface.co/models/mistralai/Mistral-7B-Instruct-v0.3"
+    
+    # Headers
+    headers = {"Authorization": f"Bearer {huggingface_token}"}
+    
+    # Payload
+    payload = {
+        "inputs": prompt,
+        "parameters": {
+            "max_new_tokens": 1000,
+            "temperature": 0.7,
+            "top_p": 0.95,
+            "top_k": 40,
+            "repetition_penalty": 1.1
+        }
+    }
+    
+    # Make the API call
+    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:
+            generated_text = result[0].get('generated_text', 'No text generated')
+            
+            # Remove the instruction part
+            content_start = generated_text.find("[/INST]")
+            if content_start != -1:
+                generated_text = generated_text[content_start + 7:].strip()
+            
+            # Split the response into main content and sources
+            parts = generated_text.split("Sources:", 1)
+            main_content = parts[0].strip()
+            sources = parts[1].strip() if len(parts) > 1 else ""
+            
+            return main_content, sources
+        else:
+            return f"Unexpected response format: {result}", ""
+    else:
+        return f"Error: API returned status code {response.status_code}", ""
 
-def chatbot_interface(message, history, temperature, top_p, repetition_penalty, use_pdf):
-    logging.debug(f"Chatbot interface called with message: {message}")
-    main_content, sources = get_response_with_search(message, temperature, top_p, repetition_penalty, use_pdf)
+def chatbot_interface(message, history):
+    main_content, sources = get_response_with_search(message)
     formatted_response = f"{main_content}\n\nSources:\n{sources}"
-    logging.debug(f"Formatted response: {formatted_response[:100]}...") # Log first 100 characters
     return formatted_response
 
-# Gradio interface
-with gr.Blocks() as demo:
-    gr.Markdown("# AI-powered Web Search and PDF Chat Assistant")
-    
-    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")
-        update_button = gr.Button("Upload PDF")
-    
-    update_output = gr.Textbox(label="Update Status")
-    update_button.click(update_vectors, inputs=[file_input, parser_dropdown], outputs=update_output)
-    
-    with gr.Row():
-        with gr.Column(scale=2):
-            chatbot = gr.Chatbot(label="Conversation")
-            msg = gr.Textbox(label="Ask a question")
-            submit_button = gr.Button("Submit")
-        with gr.Column(scale=1):
-            temperature = gr.Slider(label="Temperature", minimum=0.0, maximum=1.0, value=0.7, step=0.1)
-            top_p = gr.Slider(label="Top P", minimum=0.0, maximum=1.0, value=0.95, step=0.05)
-            repetition_penalty = gr.Slider(label="Repetition Penalty", minimum=1.0, maximum=2.0, value=1.1, step=0.1)
-            use_pdf = gr.Checkbox(label="Use PDF Documents", value=False)
-
-    def respond(message, chat_history, temperature, top_p, repetition_penalty, use_pdf):
-        bot_message = chatbot_interface(message, chat_history, temperature, top_p, repetition_penalty, use_pdf)
-        chat_history.append((message, bot_message))
-        return "", chat_history
-
-    submit_button.click(respond, inputs=[msg, chatbot, temperature, top_p, repetition_penalty, use_pdf], outputs=[msg, chatbot])
-    
-    clear_button = gr.Button("Clear Cache")
-    clear_output = gr.Textbox(label="Cache Status")
-    clear_button.click(clear_cache, inputs=[], outputs=clear_output)
+# Gradio chatbot interface
+iface = gr.ChatInterface(
+    fn=chatbot_interface,
+    title="AI-powered Web Search Assistant",
+    description="Ask questions, and I'll search the web and provide answers using the Mistral-7B-Instruct model.",
+    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"],
+    ],
+    retry_btn="Retry",
+    undo_btn="Undo",
+    clear_btn="Clear",
+)
 
 if __name__ == "__main__":
-    demo.launch()
+    iface.launch(share=True)