Update app.py

app.py

@@ -4,13 +4,23 @@ from llama_index.core import (
     VectorStoreIndex,
     SimpleDirectoryReader,
     Settings,
+    PromptTemplate,
+    QueryBundle,
 )
-from llama_index.core import PromptTemplate
 from llama_index.llms.gemini import Gemini
 from llama_index.embeddings.gemini import GeminiEmbedding
+from llama_index.core import get_response_synthesizer
+from llama_index.core.node_parser import SemanticSplitterNodeParser
+from llama_index.core.retrievers import VectorIndexRetriever
+from llama_index.core.query_engine import RetrieverQueryEngine
+from llama_index.core.query_transform import HyDEQueryTransform
+from llama_index.core.postprocessor import SentenceTransformerRerank
+from llama_index.core import load_index_from_storage
+from llama_index.core import StorageContext
+from llama_index.core.retrievers import QueryFusionRetriever
+from dotenv import load_dotenv
 import logging
 import google.generativeai as genai
-from dotenv import load_dotenv
 from pathlib import Path
 
 load_dotenv()
@@ -18,20 +28,29 @@ load_dotenv()
 # Set logging level
 logging.basicConfig(level=logging.INFO)
 
-
 # Configure Gemini Pro
 genai.configure(api_key=os.getenv("GOOGLE_API_KEY"))
 
-model_gemini_pro_vision = "gemini-1.5-flash"
+model_gemini_pro_vision = "gemini-pro-vision"
+model_gemini_pro = "gemini-pro"
+
 
 # Configure Gemini models
-Settings.llm = Gemini(model=model_gemini_pro_vision,
-                        api_key=os.getenv("GOOGLE_API_KEY"))
+Settings.llm = Gemini(model=model_gemini_pro, api_key=os.getenv("GOOGLE_API_KEY"))
 Settings.embed_model = GeminiEmbedding(
     model_name="models/embedding-001",
-                        api_key=os.getenv("GOOGLE_API_KEY")
+    api_key=os.getenv("GOOGLE_API_KEY")
+)
+
+
+# Function to create a Semantic Splitter Node Parser
+def create_semantic_splitter_node_parser():
+    """Creates a semantic splitter."""
+    return SemanticSplitterNodeParser(
+        buffer_size=1, breakpoint_percentile_threshold=95, embed_model=Settings.embed_model
     )
 
+
 def load_and_index_pdf(pdf_path):
     """Loads and index the pdf.
     
@@ -42,25 +61,69 @@ def load_and_index_pdf(pdf_path):
     index (llama_index.core.VectorStoreIndex): The vector index
     """
     try:
-        logging.info(f"Loading PDF document from: {pdf_path}")
-        documents = SimpleDirectoryReader(input_files=[pdf_path]).load_data()
-        if documents:
+         logging.info(f"Loading PDF document from: {pdf_path}")
+         documents = SimpleDirectoryReader(input_files=[pdf_path]).load_data()
+         if documents:
+            logging.info("Creating semantic splitter")
+            node_parser = create_semantic_splitter_node_parser()
+            nodes = node_parser.get_nodes_from_documents(documents)
             logging.info("Creating vector store index")
-            index = VectorStoreIndex.from_documents(documents)
+            index = VectorStoreIndex(nodes=nodes)
             return index
-        else:
+         else:
             logging.warning("No documents found in the PDF")
             return None
     except Exception as e:
         logging.error(f"Error loading and indexing PDF: {e}")
         return None
+
+
+def create_rag_pipeline(index):
+    """Creates a RAG pipeline for translation.
     
-def translate_text(french_text, index):
-    """Translates french text to Yipunu.
+     Args : 
+    index (llama_index.core.VectorStoreIndex): The vector index.
+    
+    Returns : 
+    query_engine(llama_index.core.query_engine.RetrieverQueryEngine): The query engine
+    """
+
+    logging.info("Initializing RAG Pipeline components")
+    # setup retriever
+
+    retriever = VectorStoreIndex(
+            index.nodes,
+            ).as_retriever(similarity_top_k=5)
+
+
+     # setup query transformer
+    hyde_query_transform = HyDEQueryTransform(llm=Settings.llm)
+
+      # setup reranker
+    reranker = SentenceTransformerRerank(top_n=3, model="BAAI/bge-reranker-base")
+
+    # response_synthesizer
+    response_synthesizer = get_response_synthesizer(
+            response_mode="refine",
+        )
+    
+    # setup query engine
+    query_engine = RetrieverQueryEngine(
+        retriever=retriever,
+        response_synthesizer=response_synthesizer,
+        node_postprocessors=[reranker],
+        query_transform= hyde_query_transform
+    )
+
+    logging.info("RAG Pipeline is configured.")
+    return query_engine
+
+def translate_text(french_text, query_engine):
+    """Translates french text to Yipunu using a highly optimized RAG.
     
     Args : 
     french_text (str): The french text to translate.
-    index (llama_index.core.VectorStoreIndex): The vector index.
+    query_engine (llama_index.core.query_engine.RetrieverQueryEngine): The query engine.
     
     Returns:
      (str): The yipunu translation or an error message.
@@ -76,43 +139,55 @@ def translate_text(french_text, index):
         )
 
         prompt_template = PromptTemplate(template)
-
-        query_engine = index.as_query_engine(
-            text_qa_template=prompt_template
-        )
-        response = query_engine.query(french_text)
+        query_bundle = QueryBundle(french_text, custom_prompt=prompt_template)
+        response = query_engine.query(query_bundle)
         logging.info(f"Translation Result: {response.response}")
         return response.response
     except Exception as e:
         logging.error(f"Error during translation: {e}")
         return f"Error during translation: {str(e)}"
-        
+    
+
 
 def main():
     """Main function for streamlit app."""
     
     st.title("French to Yipunu Translation App")
-
-    # PDF File Upload
-    uploaded_file = st.file_uploader("Upload a PDF file containing the Punu grammar:", type="pdf")
-
-    if uploaded_file is not None:
-            # Save file to a temporary location
-        temp_file_path = Path("temp_file.pdf")
-        with open(temp_file_path, "wb") as f:
-            f.write(uploaded_file.read())
-
-        index = load_and_index_pdf(str(temp_file_path))
+    
+    # Construct the path to the PDF in the data folder
+    default_pdf_path = Path("data/parlons_yipunu.pdf")
+    
+    # Check if the default pdf_file exists.
+    if default_pdf_path.exists():
+        index = load_and_index_pdf(str(default_pdf_path))
         if index:
+            query_engine = create_rag_pipeline(index)
             french_text = st.text_area("Enter French Text:", "Ni vosi yipunu")
             if st.button("Translate"):
-                 translation = translate_text(french_text, index)
-                 st.success(f"Yipunu Translation: {translation}")
-
-        # Clean up temp files
-        os.remove(temp_file_path)
+                translation = translate_text(french_text, query_engine)
+                st.success(f"Yipunu Translation: {translation}")
     else:
-        st.info("Please upload a pdf containing the punu grammar.")
+            # PDF File Upload
+        uploaded_file = st.file_uploader("Upload a PDF file containing the Punu grammar:", type="pdf")
+        if uploaded_file is not None:
+            # Save file to a temporary location
+            temp_file_path = Path("temp_file.pdf")
+            with open(temp_file_path, "wb") as f:
+                f.write(uploaded_file.read())
+
+            index = load_and_index_pdf(str(temp_file_path))
+            if index:
+                query_engine = create_rag_pipeline(index)
+                french_text = st.text_area("Enter French Text:", "Ni vosi yipunu")
+                if st.button("Translate"):
+                    translation = translate_text(french_text, query_engine)
+                    st.success(f"Yipunu Translation: {translation}")
+
+            # Clean up temp files
+            os.remove(temp_file_path)
+        else:
+            st.info("Please upload a pdf containing the punu grammar.")
+    
 
 
 if __name__ == "__main__":