Update app.py

app.py

@@ -1,30 +1,129 @@
-import gradio as gr
-from autollm import AutoQueryEngine
 import os
+import dotenv
+import gradio as gr
+import lancedb
+import logging
+from langchain.embeddings.cohere import CohereEmbeddings
+from langchain.llms import Cohere
+from langchain.prompts import PromptTemplate
+from langchain.chains import RetrievalQA
+from langchain.vectorstores import LanceDB
+from langchain.document_loaders import TextLoader
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain_community.document_loaders import PyPDFLoader
+import argostranslate.package
+import argostranslate.translate
 import spaces
 
+
+# Configuration Management
+dotenv.load_dotenv(".env")
+DB_PATH = "/tmp/lancedb"
+
+COHERE_MODEL_NAME = "multilingual-22-12"
+LANGUAGE_ISO_CODES = {
+    "English": "en",
+    "Hindi": "hi",
+    "Turkish": "tr",
+    "French": "fr",
+}
+
+# Logging Configuration
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
 @spaces.GPU
-def query_engine(llm_model, document, query):
-    api_key = os.getenv("HUGGINGFACE_API_KEY")
-    llm_api_base = "https://api-inference.huggingface.co/models/"
+def initialize_documents_and_embeddings(input_file_path):
+    file_extension = os.path.splitext(input_file_path)[1]
+    if file_extension == '.txt':
+        logger.info("txt file processing")
+        # Handle text file
+        loader = TextLoader(input_file_path)
+        documents = loader.load()
+        text_splitter = RecursiveCharacterTextSplitter(chunk_size=300, chunk_overlap=50)
+        texts = text_splitter.split_documents(documents)
+    elif file_extension == '.pdf':
+        logger.info("pdf file processing")
+        # Handle PDF file
+        loader = PyPDFLoader(input_file_path)
+        texts = loader.load_and_split()
+        text_splitter = RecursiveCharacterTextSplitter(chunk_size=300, chunk_overlap=50)
+        texts = text_splitter.split_documents(texts)
+    else:
+        raise ValueError("Unsupported file type. Supported files are .txt and .pdf only.")
+    
+    embeddings = CohereEmbeddings(model=COHERE_MODEL_NAME)
+    return texts, embeddings
+
+# Database Initialization
+def initialize_database(texts, embeddings):
+    db = lancedb.connect(DB_PATH)
+    table = db.create_table(
+        "multiling-rag",
+        data=[{"vector": embeddings.embed_query("Hello World"), "text": "Hello World", "id": "1"}],
+        mode="overwrite",
+    )
+    return LanceDB.from_documents(texts, embeddings, connection=table)
 
-    content = str(document)
+# Translation Function
+def translate_text(text, from_code, to_code):
+    try:
+        argostranslate.package.update_package_index()
+        available_packages = argostranslate.package.get_available_packages()
+        package_to_install = next(filter(lambda x: x.from_code == from_code and x.to_code == to_code, available_packages))
+        argostranslate.package.install_from_path(package_to_install.download())
+        return argostranslate.translate.translate(text, from_code, to_code)
+    except Exception as e:
+        logger.error(f"Error in translate_text: {str(e)}")
+        return "Translation error"
 
-    query_engine = AutoQueryEngine.from_defaults(
-        documents=[content],
-        llm_model=llm_model,
-        llm_api_base=llm_api_base,
+
+prompt_template = """Text: {context}
+
+Question: {question}
+
+Answer the question based on the text provided. If the text doesn't contain the answer, reply that the answer is not available."""
+PROMPT = PromptTemplate(
+    template=prompt_template, input_variables=["context", "question"])
+
+# Question Answering Function
+def answer_question(question, input_language, output_language, db):
+    try:
+        input_lang_code = LANGUAGE_ISO_CODES[input_language]
+        output_lang_code = LANGUAGE_ISO_CODES[output_language]
+        
+        question_in_english = translate_text(question, from_code=input_lang_code, to_code="en") if input_language != "English" else question
+        prompt = PromptTemplate(template=prompt_template, input_variables=["context", "question"])
+        qa = RetrievalQA.from_chain_type(llm=Cohere(model="command", temperature=0), chain_type="stuff", retriever=db.as_retriever(), chain_type_kwargs={"prompt": prompt}, return_source_documents=True)
+        
+        answer = qa({"query": question_in_english})
+        result_in_english = answer["result"].replace("\n", "").replace("Answer:", "")
+        
+        return translate_text(result_in_english, from_code="en", to_code=output_lang_code) if output_language != "English" else result_in_english
+    except Exception as e:
+        logger.error(f"Error in answer_question: {str(e)}")
+        return "An error occurred while processing your question. Please try again."
+
+def setup_gradio_interface(db):
+    return gr.Interface(
+        fn=lambda question, input_language, output_language: answer_question(question, input_language, output_language, db),
+        inputs=[
+            gr.Textbox(lines=2, placeholder="Type your question here..."),
+            gr.Dropdown(list(LANGUAGE_ISO_CODES.keys()), label="Input Language"),
+            gr.Dropdown(list(LANGUAGE_ISO_CODES.keys()), label="Output Language")
+        ],
+        outputs="text",
+        title="Multilingual Chatbot",
+        description="Ask any question in your chosen language and get an answer in the language of your choice."
     )
-    response = query_engine.query(query)
-    return response
-
-interface = gr.Blocks()
-with interface:
-    gr.Markdown("# AutoQueryEngine Interface")
-    llm_model = gr.Textbox(label="LLM Model", value="mistralai/Mixtral-8x7B-Instruct-v0.1")
-    document = gr.File(label="Documents")
-    query = gr.Textbox(label="Query")
-    output = gr.Textbox(label="Output")
-    query_btn = gr.Button("Query")
-    query_btn.click(fn=query_engine, inputs=[llm_model, document, query], outputs=output)
-interface.launch()
+
+# Main Function
+def main():
+    INPUT_FILE_PATH = "healthy-diet-fact-sheet-394.pdf"
+    texts, embeddings = initialize_documents_and_embeddings(INPUT_FILE_PATH)
+    db = initialize_database(texts, embeddings)
+    iface = setup_gradio_interface(db)
+    iface.launch(share=True, debug=True)
+
+if __name__ == "__main__":
+    main()