Sissa2 / app.py
christian1984's picture
Update app.py
7373789 verified
raw
history blame
2.86 kB
import os
import streamlit as st
import requests
from langchain_community.document_loaders import PDFPlumberLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_core.vectorstores import InMemoryVectorStore
from langchain_core.prompts import ChatPromptTemplate
HUGGINGFACE_API_KEY = os.getenv("HUGGINGFACE_API_KEY")
HF_API_URL = "https://api-inference.huggingface.co/models/deepseek-ai/deepseek-llm-7b"
HEADERS = {"Authorization": f"Bearer {HUGGINGFACE_API_KEY}"}
pdfs_directory = "./pdfs/"
template = """
You are an assistant for question-answering tasks. Use the following pieces of retrieved context to answer the question. If you don't know the answer, just say that you don't know. Use three sentences maximum and keep the answer concise.
Question: {question}
Context: {context}
Answer:
"""
vector_store = InMemoryVectorStore()
def upload_pdf(file):
"""Save uploaded PDF file"""
with open(pdfs_directory + file.name, "wb") as f:
f.write(file.getbuffer())
def load_pdf(file_path):
"""Load text from the PDF"""
loader = PDFPlumberLoader(file_path)
documents = loader.load()
return documents
def split_text(documents):
"""Split text into smaller chunks"""
text_splitter = RecursiveCharacterTextSplitter(
chunk_size=1000,
chunk_overlap=200,
add_start_index=True,
)
return text_splitter.split_documents(documents)
def index_docs(documents):
vector_store.add_documents(documents)
def retrieve_docs(query):
"""Retrieve similar documents"""
return vector_store.similarity_search(query)
def query_huggingface_api(prompt):
"""Send query to DeepSeek R1 model on Hugging Face"""
payload = {"inputs": prompt}
response = requests.post(HF_API_URL, headers=HEADERS, json=payload)
if response.status_code == 200:
return response.json()[0]["generated_text"]
else:
return "Error: Unable to process request."
def answer_question(question, documents):
"""Generate answer using DeepSeek R1 API"""
context = "\n\n".join([doc.page_content for doc in documents])
prompt = ChatPromptTemplate.from_template(template).format(question=question, context=context)
return query_huggingface_api(prompt)
st.title("PDF-based RAG Chatbot")
uploaded_file = st.file_uploader("Upload PDF", type="pdf", accept_multiple_files=False)
if uploaded_file:
upload_pdf(uploaded_file)
documents = load_pdf(pdfs_directory + uploaded_file.name)
chunked_documents = split_text(documents)
index_docs(chunked_documents)
question = st.chat_input("Ask a question based on the document")
if question:
st.chat_message("user").write(question)
related_documents = retrieve_docs(question)
answer = answer_question(question, related_documents)
st.chat_message("assistant").write(answer)