Spaces:
Sleeping
Sleeping
| """ | |
| Gradio UI for Mistral 7B with RAG | |
| """ | |
| import os | |
| from typing import List | |
| import gradio as gr | |
| from langchain_core.runnables.base import RunnableSequence | |
| import numpy as np | |
| from confluence_rag import generate_rag_chain, load_pdf, store_vector, load_multiple_pdf | |
| def initialize_chain(file: gr.File) -> RunnableSequence: | |
| """ | |
| Initializes the chain with the given file. | |
| If no file is provided, the llm is used without RAG. | |
| Args: | |
| file (gr.File): file to initialize the chain with | |
| Returns: | |
| RunnableSequence: the chain | |
| """ | |
| if file is None: | |
| return generate_rag_chain() | |
| if len(file) == 1: | |
| pdf = load_pdf(file[0].name) | |
| else: | |
| pdf = load_multiple_pdf([f.name for f in file]) | |
| retriever = store_vector(pdf) | |
| return generate_rag_chain(retriever) | |
| def invoke_chain(message: str, history: List[str], file: gr.File = None) -> str: | |
| """ | |
| Invokes the chain with the given message and updates the chain if a new file is provided. | |
| Args: | |
| message (str): message to invoke the chain with | |
| history (List[str]): history of messages | |
| file (gr.File, optional): file to update the chain with. Defaults to None. | |
| Returns: | |
| str: the response of the chain | |
| """ | |
| # Check if file is provided and exists | |
| if file is not None and not np.all([os.path.exists(f.name) for f in file]) or len(file) == 0: | |
| return "Error: File not found." | |
| if file is not None and not np.all([f.name.endswith(".pdf") for f in file]): | |
| return "Error: File is not a pdf." | |
| chain = initialize_chain(file) | |
| return chain.invoke(message) | |
| def create_demo() -> gr.Interface: | |
| """ | |
| Creates and returns a Gradio Chat Interface. | |
| Returns: | |
| gr.Interface: the Gradio Chat Interface | |
| """ | |
| return gr.ChatInterface( | |
| invoke_chain, | |
| additional_inputs=[gr.File(label="File", file_count='multiple')], | |
| title="Mistral 7B with RAG", | |
| description="Ask questions to Mistral about your pdf document.", | |
| theme="soft", | |
| ) | |
| if __name__ == "__main__": | |
| demo = create_demo() | |
| demo.launch() |