Update app.py

app.py

@@ -1,115 +1,96 @@
 import gradio as gr
+from llama_index.core import StorageContext, load_index_from_storage, VectorStoreIndex, SimpleDirectoryReader, ChatPromptTemplate
+from llama_index.llms.huggingface import HuggingFaceInferenceAPI
+from dotenv import load_dotenv
+from llama_index.embeddings.huggingface import HuggingFaceEmbedding
+from llama_index.core import Settings
 import os
-from langchain_community.embeddings import HuggingFaceEmbeddings
-from langchain_community.vectorstores import Chroma
-from langchain.text_splitter import RecursiveCharacterTextSplitter
-from langchain_community.document_loaders import TextLoader, PyPDFLoader
-from langchain.chains import RetrievalQA
-from langchain_community.llms import HuggingFaceHub
-import tempfile
-import shutil
-from langchain.prompts import PromptTemplate
- 
-# Define a proper prompt template
-prompt_template = """Use the following pieces of context to answer the question at the end. If you don't know the answer, just say that you don't know, don't try to make up an answer.
-{context}
-Question: {question}
-Answer:"""
-PROMPT = PromptTemplate(
-    template=prompt_template, input_variables=["context", "question"]
-)
- 
+import base64
+
 # Load environment variables
-TOKEN = os.getenv("HF_TOKEN")
-os.environ["HUGGINGFACEHUB_API_TOKEN"] = TOKEN
- 
-# Initialize LangChain components
-embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
- 
-# Create vector store
-vectorstore = Chroma(persist_directory="./chroma_db", embedding_function=embeddings)
- 
-# Initialize LLM
-llm = HuggingFaceHub(
-    repo_id="meta-llama/Meta-Llama-3.1-405B-Instruct-FP8",
-    model_kwargs={"temperature": 0.7, "max_length": 512}
+load_dotenv()
+
+# Configure the Llama index settings
+Settings.llm = HuggingFaceInferenceAPI(
+    model_name="nltpt/Llama-3.2-3B-Instruct",
+    tokenizer_name="nltpt/Llama-3.2-3B-Instruct",
+    context_window=3000,
+    token=os.getenv("HF_TOKEN"),
+    max_new_tokens=512,
+    generate_kwargs={"temperature": 0.1},
+)
+Settings.embed_model = HuggingFaceEmbedding(
+    model_name="BAAI/bge-small-en-v1.5"
 )
- 
-# Create RetrievalQA chain
-qa_chain = RetrievalQA.from_chain_type(
-    llm=llm,
-    chain_type="stuff",
-    retriever=vectorstore.as_retriever(),
-    return_source_documents=True,
-    chain_type_kwargs={"prompt": PROMPT}
+
+# Define the directory for persistent storage and data
+PERSIST_DIR = "./db"
+DATA_DIR = "data"
+
+# Ensure data directory exists
+os.makedirs(DATA_DIR, exist_ok=True)
+os.makedirs(PERSIST_DIR, exist_ok=True)
+
+def displayPDF(file):
+    with open(file, "rb") as f:
+        base64_pdf = base64.b64encode(f.read()).decode('utf-8')
+    pdf_display = f'<iframe src="data:application/pdf;base64,{base64_pdf}" width="100%" height="600" type="application/pdf"></iframe>'
+    return pdf_display
+
+def data_ingestion(files):
+    for file in files:
+        filepath = os.path.join(DATA_DIR, file.name)
+        with open(filepath, "wb") as f:
+            f.write(file.getbuffer())
+    documents = SimpleDirectoryReader(DATA_DIR).load_data()
+    storage_context = StorageContext.from_defaults()
+    index = VectorStoreIndex.from_documents(documents)
+    index.storage_context.persist(persist_dir=PERSIST_DIR)
+
+def handle_query(query):
+    storage_context = StorageContext.from_defaults(persist_dir=PERSIST_DIR)
+    index = load_index_from_storage(storage_context)
+    chat_text_qa_msgs = [
+        (
+            "user",
+            """You are a Q&A assistant. Your main goal is to provide answers as accurately as possible, based on the context of the document provided. If the question does not match the context or is outside the scope of the document, advise the user to ask questions that are relevant to the document.
+            Context:
+            {context_str}
+            Question:
+            {query_str}
+            """
+        )
+    ]
+    
+    text_qa_template = ChatPromptTemplate.from_messages(chat_text_qa_msgs)
+    
+    query_engine = index.as_query_engine(text_qa_template=text_qa_template)
+    answer = query_engine.query(query)
+    
+    if hasattr(answer, 'response'):
+        return answer.response
+    elif isinstance(answer, dict) and 'response' in answer:
+        return answer['response']
+    else:
+        return "Sorry, I couldn't find an answer."
+
+# Gradio app setup
+def gradio_app(files, user_query):
+    if files:
+        data_ingestion(files)  # Process PDFs after they are uploaded
+        response = handle_query(user_query)
+        return response
+    return "Please upload at least one PDF file."
+
+interface = gr.Interface(
+    fn=gradio_app,
+    inputs=[
+        gr.File(label="Upload PDF Files", type="file", file_count="multiple"),
+        gr.Textbox(label="Ask me anything about the content of the PDF(s):")
+    ],
+    outputs="text",
+    title="(PDF) Information and Inference🗞️",
+    description="Retrieval-Augmented Generation. Start chat ...🚀"
 )
- 
-def process_uploaded_file(file):
-    if file is None:
-        return "No file uploaded."
-    try:
-        # Create a temporary directory
-        with tempfile.TemporaryDirectory() as temp_dir:
-            # Create a path for the temporary file
-            temp_file_path = os.path.join(temp_dir, os.path.basename(file.name))
-            # Save the uploaded file to the temporary path
-            with open(temp_file_path, 'wb') as temp_file:
-                temp_file.write(file.read())
-            # Determine file type and load accordingly
-            if file.name.endswith('.pdf'):
-                loader = PyPDFLoader(temp_file_path)
-            else:
-                loader = TextLoader(temp_file_path)
- 
-            documents = loader.load()
-            text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
-            texts = text_splitter.split_documents(documents)
- 
-            # Update the vector store with new documents
-            vectorstore.add_documents(texts)
-            vectorstore.persist()
- 
-        return f"File processed and added to the knowledge base. {len(texts)} chunks created."
-    except Exception as e:
-        return f"An error occurred while processing the file: {str(e)}"
- 
-def respond(message, history, system_message, max_tokens, temperature, top_p):
-    full_prompt = f"{system_message}\n\nHuman: {message}"
-    # Use the RetrievalQA chain to get the answer
-    result = qa_chain({"query": full_prompt})
-    answer = result['result']
-    # Return only the answer
-    yield answer
- 
-with gr.Blocks() as demo:
-    gr.Markdown("# RAG Chatbot with Content Upload")
-    with gr.Row():
-        with gr.Column(scale=3):
-            chatbot = gr.Chatbot()
-            msg = gr.Textbox()
-            clear = gr.Button("Clear")
-        with gr.Column(scale=1):
-            file_upload = gr.File(label="Upload Content for RAG (TXT or PDF)")
-            upload_button = gr.Button("Process Uploaded File")
-            system_message = gr.Textbox(value="You are a friendly Chatbot.", label="System message")
-            max_tokens = gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens")
-            temperature = gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature")
-            top_p = gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-p (nucleus sampling)")
- 
-    def user(user_message, history):
-        return "", history + [[user_message, None]]
- 
-    def bot(history, system_message, max_tokens, temperature, top_p):
-        user_message = history[-1][0]
-        bot_message = next(respond(user_message, history[:-1], system_message, max_tokens, temperature, top_p))
-        history[-1][1] = bot_message
-        return history
- 
-    msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False).then(
-        bot, [chatbot, system_message, max_tokens, temperature, top_p], chatbot
-    )
-    clear.click(lambda: None, None, chatbot, queue=False)
-    upload_button.click(process_uploaded_file, inputs=[file_upload], outputs=[gr.Textbox()])
- 
-if __name__ == "__main__":
-    demo.launch()
+
+interface.launch()