Spaces:
Sleeping
Sleeping
| import warnings | |
| warnings.filterwarnings("ignore") | |
| import uuid | |
| import gradio as gr | |
| import os | |
| from langchain_huggingface import HuggingFaceEndpoint, ChatHuggingFace | |
| from langchain_text_splitters import RecursiveCharacterTextSplitter | |
| from langchain_community.vectorstores import Chroma | |
| from langchain_community.document_loaders import PyPDFLoader | |
| from langchain.chains import RetrievalQA | |
| from langchain_community.embeddings import HuggingFaceEmbeddings | |
| # ------------------------- | |
| # LLM (Hugging Face Router API) | |
| # ------------------------- | |
| repo_id = "Qwen/Qwen2.5-7B-Instruct" | |
| llm_endpoint = HuggingFaceEndpoint( | |
| repo_id=repo_id, | |
| task="text-generation", | |
| max_new_tokens=512, | |
| temperature=0.1, | |
| huggingfacehub_api_token=os.getenv("HF_TOKEN"), | |
| provider="hf-inference", # forces router.huggingface.co | |
| model_kwargs={"wait_for_model": True}, | |
| ) | |
| llm = ChatHuggingFace(llm=llm_endpoint) | |
| # ------------------------- | |
| # Document loader | |
| # ------------------------- | |
| def document_loader(file_path: str): | |
| loader = PyPDFLoader(file_path) | |
| return loader.load() | |
| # ------------------------- | |
| # Text splitter | |
| # ------------------------- | |
| def text_splitter(data): | |
| splitter = RecursiveCharacterTextSplitter( | |
| chunk_size=1000, | |
| chunk_overlap=50, | |
| length_function=len, | |
| ) | |
| return splitter.split_documents(data) | |
| # ------------------------- | |
| # Embedding model | |
| # ------------------------- | |
| def embedding_model(): | |
| return HuggingFaceEmbeddings( | |
| model_name="sentence-transformers/all-MiniLM-L6-v2" | |
| ) | |
| # ------------------------- | |
| # Vector DB | |
| # ------------------------- | |
| def vector_database(chunks): | |
| chunks = [ | |
| c for c in chunks | |
| if c.page_content and c.page_content.strip() and len(c.page_content.strip()) > 10 | |
| ] | |
| vectordb = Chroma.from_documents( | |
| documents=chunks, | |
| embedding=embedding_model(), | |
| collection_name=f"pdf_chat_{uuid.uuid4().hex}", | |
| ) | |
| return vectordb | |
| # ------------------------- | |
| # Retriever | |
| # ------------------------- | |
| def build_retriever(file_path: str): | |
| docs = document_loader(file_path) | |
| chunks = text_splitter(docs) | |
| vectordb = vector_database(chunks) | |
| return vectordb.as_retriever() | |
| # ------------------------- | |
| # QA Chain | |
| # ------------------------- | |
| def retriever_qa(file, query): | |
| retr = build_retriever(file) | |
| qa = RetrievalQA.from_chain_type( | |
| llm=llm, | |
| chain_type="stuff", | |
| retriever=retr, | |
| return_source_documents=False, | |
| ) | |
| result = qa.invoke({"query": query}) | |
| return result["result"] | |
| # ------------------------- | |
| # Gradio interface | |
| # ------------------------- | |
| rag_application = gr.Interface( | |
| fn=retriever_qa, | |
| allow_flagging="never", | |
| inputs=[ | |
| gr.File( | |
| label="Upload PDF File", | |
| file_count="single", | |
| file_types=[".pdf"], | |
| type="filepath", | |
| ), | |
| gr.Textbox( | |
| label="Input Query", | |
| lines=2, | |
| placeholder="Type your question here...", | |
| ), | |
| ], | |
| outputs=gr.Textbox(label="Output"), | |
| title="PDF QA Chatbot", | |
| description=( | |
| "Upload a PDF document and ask any question. " | |
| "The chatbot will try to answer using the provided document." | |
| ), | |
| ) | |
| if __name__ == "__main__": | |
| rag_application.launch() |