Sowmya135/RetrievalAugmentedGenerator

1import gradio as gr2from langchain_community.document_loaders import PyPDFLoader, TextLoader, Docx2txtLoader3from langchain_text_splitters import RecursiveCharacterTextSplitter4from langchain_huggingface import HuggingFaceEmbeddings5from langchain_community.vectorstores import FAISS6from langchain_community.llms import HuggingFacePipeline7from langchain_core.prompts import PromptTemplate8from langchain_core.output_parsers import StrOutputParser9from langchain_core.runnables import RunnablePassthrough10from transformers import pipeline11 12vectorstore = None13 14 15def load_documents(files):16    documents = []17 18    for file in files:19        path = file.name20 21        if path.endswith(".pdf"):22            loader = PyPDFLoader(path)23        elif path.endswith(".txt"):24            loader = TextLoader(path)25        elif path.endswith(".docx"):26            loader = Docx2txtLoader(path)27        else:28            continue29 30        documents.extend(loader.load())31 32    return documents33 34 35def format_docs(docs):36    return "\n\n".join(doc.page_content for doc in docs)37 38 39def build_vectorstore(files):40    docs = load_documents(files)41 42    splitter = RecursiveCharacterTextSplitter(43        chunk_size=500,44        chunk_overlap=5045    )46    splits = splitter.split_documents(docs)47 48    embeddings = HuggingFaceEmbeddings(49        model_name="sentence-transformers/all-MiniLM-L6-v2"50    )51 52    return FAISS.from_documents(splits, embeddings)53 54 55def build_chain(vs):56    retriever = vs.as_retriever(search_kwargs={"k": 4})57 58    llm = HuggingFacePipeline(59        pipeline=pipeline(60            "text2text-generation",61            model="google/flan-t5-base",62            max_new_tokens=25663        )64    )65 66    prompt = PromptTemplate.from_template(67        """Answer the question using ONLY the context below.68If the answer is not in the context, say "I don't know."69 70Context:71{context}72 73Question:74{question}75 76Answer:77"""78    )79 80    return (81        {82            "context": retriever | format_docs,83            "question": RunnablePassthrough()84        }85        | prompt86        | llm87        | StrOutputParser()88    )89 90 91def chat(files, question):92    global vectorstore93 94    if not files:95        return "Please upload at least one document."96 97    if vectorstore is None:98        vectorstore = build_vectorstore(files)99 100    chain = build_chain(vectorstore)101    return chain.invoke(question)102 103 104iface = gr.Interface(105    fn=chat,106    inputs=[107        gr.File(file_types=[".pdf", ".txt", ".docx"], file_count="multiple"),108        gr.Textbox(label="Ask a question")109    ],110    outputs=gr.Textbox(label="Answer", lines=10),111    title="📄 Doc Query RAG"112)113 114iface.launch()115 116