Update gradio_app.py

gradio_app.py

@@ -6,29 +6,17 @@ from llama_index.embeddings import HuggingFaceEmbedding
 from llama_index import ServiceContext
 from dotenv import load_dotenv
 import os
+import base64
 import tempfile
 from pathlib import Path
-import logging
-
-# Set up logging
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
 
 # Load environment variables
 load_dotenv()
 
-# 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)
-
 # Configure the Llama index settings
 llm = HuggingFaceInferenceAPI(
-    model_name="google/gemma-1.1-7b-it",
-    tokenizer_name="google/gemma-1.1-7b-it",
+    model_name="google/gemini-1.1-7b-it",
+    tokenizer_name="google/gemini-1.1-7b-it",
     context_window=3000,
     token=os.getenv("HF_TOKEN"),
     max_new_tokens=512,
@@ -44,57 +32,58 @@ service_context = ServiceContext.from_defaults(
     embed_model=embed_model
 )
 
+# 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):
+    base64_pdf = base64.b64encode(file.read()).decode('utf-8')
+    return f'<iframe src="data:application/pdf;base64,{base64_pdf}" width="100%" height="600" type="application/pdf"></iframe>'
+
 def data_ingestion():
-    try:
-        documents = SimpleDirectoryReader(DATA_DIR).load_data()
-        storage_context = StorageContext.from_defaults()
-        index = VectorStoreIndex.from_documents(documents, service_context=service_context)
-        index.storage_context.persist(persist_dir=PERSIST_DIR)
-        return True
-    except Exception as e:
-        logger.error(f"Error during data ingestion: {str(e)}")
-        return False
+    documents = SimpleDirectoryReader(DATA_DIR).load_data()
+    storage_context = StorageContext.from_defaults()
+    index = VectorStoreIndex.from_documents(documents, service_context=service_context)
+    index.storage_context.persist(persist_dir=PERSIST_DIR)
 
 def handle_query(query):
-    try:
-        storage_context = StorageContext.from_defaults(persist_dir=PERSIST_DIR)
-        index = load_index_from_storage(storage_context, service_context=service_context)
-        chat_text_qa_msgs = [
-        (
-            "user",
-            """You are a Q&A assistant named EazyPeazy, For all other inquiries, your main goal is to provide answers as accurately as possible, based on the instructions and context you have been given. If a question does not match the provided context or is outside the scope of the document, kindly advise the user to ask questions within the context of 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."
-    except Exception as e:
-        logger.error(f"Error handling query: {str(e)}")
-        return "An error occurred while processing your query. Please try again."
+    storage_context = StorageContext.from_defaults(persist_dir=PERSIST_DIR)
+    index = load_index_from_storage(storage_context, service_context=service_context)
+    chat_text_qa_msgs = [
+    (
+        "user",
+        """You are a Q&A assistant named CHATTO, created by Suriya. You have a specific response programmed for when users specifically ask about your creator, Suriya. The response is: "I was created by Suriya, an enthusiast in Artificial Intelligence. He is dedicated to solving complex problems and delivering innovative solutions. With a strong focus on machine learning, deep learning, Python, generative AI, NLP, and computer vision, Suriya is passionate about pushing the boundaries of AI to explore new possibilities." For all other inquiries, your main goal is to provide answers as accurately as possible, based on the instructions and context you have been given. If a question does not match the provided context or is outside the scope of the document, kindly advise the user to ask questions within the context of 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."
 
 def process_file(file):
     if file is None:
         return "Please upload a PDF file."
-    
+
     try:
-        temp_dir = tempfile.mkdtemp()
-        temp_path = Path(temp_dir) / file.name
-        
-        # Write the file content to the temporary file
-        temp_path.write_bytes(file.read())
+        with tempfile.NamedTemporaryFile(delete=False) as temp_file:
+            temp_file.write(file.read())
+            temp_path = Path(temp_file.name)
         
         # Copy the file to the DATA_DIR
         dest_path = Path(DATA_DIR) / file.name
@@ -102,18 +91,15 @@ def process_file(file):
         temp_path.replace(dest_path)
         
         # Process the uploaded PDF
-        if data_ingestion():
-            return f"PDF '{file.name}' processed successfully. You can now ask questions about its content."
-        else:
-            return f"Failed to process the PDF '{file.name}'. Please try uploading again."
+        data_ingestion()
+        return f"PDF '{file.name}' processed successfully. You can now ask questions about its content.", displayPDF(file)
     except Exception as e:
-        logger.error(f"Error processing file: {str(e)}")
-        return f"An error occurred while processing the file: {str(e)}"
+        return f"An error occurred while processing the file: {str(e)}", None
 
 def chat_function(message, history):
     response = handle_query(message)
     history.append((message, response))
-    return history
+    return history, history
 
 with gr.Blocks() as demo:
     gr.Markdown("# (PDF) Information and Inference🗞️")
@@ -122,6 +108,7 @@ with gr.Blocks() as demo:
     with gr.Row():
         with gr.Column(scale=1):
             file_output = gr.Textbox(label="Upload Status")
+            file_display = gr.HTML()
             upload_button = gr.UploadButton("Upload PDF", file_types=[".pdf"])
         
         with gr.Column(scale=2):
@@ -129,9 +116,9 @@ with gr.Blocks() as demo:
             msg = gr.Textbox(label="Ask me anything about the content of the PDF:")
             clear = gr.Button("Clear")
     
-    upload_button.upload(process_file, upload_button, file_output)
-    msg.submit(chat_function, [msg, chatbot], chatbot)
+    upload_button.upload(process_file, upload_button, [file_output, file_display])
+    msg.submit(chat_function, [msg, chatbot], [chatbot, chatbot])
     clear.click(lambda: None, None, chatbot, queue=False)
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()