Spaces:
Runtime error
Runtime error
File size: 2,027 Bytes
5588cb0 15d650d 5588cb0 15d650d 5588cb0 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 |
import os
import requests
import streamlit as st
from langchain.callbacks.base import BaseCallbackHandler
@st.cache_resource(ttl="1h")
def upload_data(uploaded_files):
files = {"file": uploaded_files}
with st.spinner("Uploading PDF..."):
response = requests.post(
"http://127.0.0.1:8000/api/upload", files=files
)
if response.status_code == 200:
st.success(
f'{response.json()["message"][0]} Vector Store created successfully!'
)
st.session_state.uploaded_pdf = True
else:
st.error("Failed to upload PDF!")
class StreamHandler(BaseCallbackHandler):
def __init__(
self, container: st.delta_generator.DeltaGenerator, initial_text: str = ""
):
self.container = container
self.text = initial_text
self.run_id_ignore_token = None
def on_llm_start(self, serialized: dict, prompts: list, **kwargs):
# Workaround to prevent showing the rephrased question as output
if prompts[0].startswith("Human"):
self.run_id_ignore_token = kwargs.get("run_id")
def on_llm_new_token(self, token: str, **kwargs) -> None:
if self.run_id_ignore_token == kwargs.get("run_id", False):
return
self.text += token
self.container.markdown(self.text)
class PrintRetrievalHandler(BaseCallbackHandler):
def __init__(self, container):
self.status = container.status("**Context Retrieval**")
def on_retriever_start(self, serialized: dict, query: str, **kwargs):
self.status.write(f"**Question:** {query}")
self.status.update(label=f"**Context Retrieval:** {query}")
def on_retriever_end(self, documents, **kwargs):
for idx, doc in enumerate(documents):
source = os.path.basename(doc.metadata["source"])
self.status.write(f"**Document {idx} from {source}**")
self.status.markdown(doc.page_content)
self.status.update(state="complete")
|