Spaces:
Runtime error
Runtime error
File size: 5,517 Bytes
51fe9d2 |
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 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 |
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.vectorstores.faiss import FAISS
from langchain import OpenAI
from langchain.chains.qa_with_sources import load_qa_with_sources_chain
from langchain.embeddings.openai import OpenAIEmbeddings
from langchain.llms import OpenAI
from langchain.docstore.document import Document
from langchain.vectorstores import FAISS, VectorStore
import docx2txt
from typing import List, Dict, Any, Union, Text, Tuple
import re
from io import BytesIO
import streamlit as st
from .prompts import STUFF_PROMPT
from pypdf import PdfReader
from openai.error import AuthenticationError
class HashDocument(Document):
"""A document that uses the page content as the hash."""
def __hash__(self):
content = self.page_content + "".join(self.metadata[k] for k in self.metadata.keys())
return hash(content)
@st.cache_data
def parse_docx(file: BytesIO) -> str:
text = docx2txt.process(file)
# Remove multiple newlines
text = re.sub(r"\n\s*\n", "\n\n", text)
return text
@st.cache_data
def parse_pdf(file: BytesIO) -> List[str]:
pdf = PdfReader(file)
output = []
for page in pdf.pages:
text = page.extract_text()
# Merge hyphenated words
text = re.sub(r"(\w+)-\n(\w+)", r"\1\2", text)
# Fix newlines in the middle of sentences
text = re.sub(r"(?<!\n\s)\n(?!\s\n)", " ", text.strip())
# Remove multiple newlines
text = re.sub(r"\n\s*\n", "\n\n", text)
output.append(text)
return output
@st.cache_data
def parse_txt(file: BytesIO) -> str:
text = file.read().decode("utf-8")
# Remove multiple newlines
text = re.sub(r"\n\s*\n", "\n\n", text)
return text
@st.cache_data
def text_to_docs(text: Union[Text, Tuple[Text]]) -> List[Document]:
"""
Converts a string or frozenset of strings to a list of Documents
with metadata.
"""
if isinstance(text, str):
# Take a single string as one page
text = tuple([text])
elif isinstance(text, tuple):
# map each page into a document instance
page_docs = [HashDocument(page_content=page) for page in text]
# Add page numbers as metadata
for i, doc in enumerate(page_docs):
doc.metadata["page"] = i + 1
# Split pages into chunks
doc_chunks = []
# text splitter to split the text into chunks
text_splitter = RecursiveCharacterTextSplitter(
chunk_size=800,
separators=["\n\n", "\n", ".", "!", "?", ",", " ", ""],
chunk_overlap=20, # minimal overlap to capture sematic overlap across chunks
)
for doc in page_docs:
chunks = text_splitter.split_text(doc.page_content)
for i, chunk in enumerate(chunks):
# Create a new document for each individual chunk
doc = HashDocument(
page_content=chunk, metadata={"page": doc.metadata["page"], "chunk": i}
)
# Add sources a metadata
doc.metadata["source"] = f"{doc.metadata['page']}-{doc.metadata['chunk']}"
doc_chunks.append(doc)
return doc_chunks
else:
raise ValueError("Text must be either a string or a list of strings. Got: {type(text)}")
@st.cache_data
def embed_docs(_docs: Tuple[Document]) -> VectorStore:
"""Embeds a list of Documents and returns a FAISS index"""
docs = _docs
if not st.session_state.get("OPENAI_API_KEY"):
raise AuthenticationError(
"Enter your OpenAI API key in the sidebar. You can get a key at https://platform.openai.com/account/api-keys."
)
else:
# Embed the chunks
embeddings = OpenAIEmbeddings(openai_api_key=st.session_state.get("OPENAI_API_KEY"))
index = FAISS.from_documents(list(docs), embeddings)
return index
@st.cache_data
def search_docs(_index: VectorStore, query: str) -> List[Document]:
"""Searches a FAISS index for similar chunks to the query
and returns a list of Documents."""
# Search for similar chunks
docs = _index.similarity_search(query, k=5)
return docs
@st.cache_data
def get_answer(_docs: List[Document], query: str) -> Dict[str, Any]:
"""Gets an answer to a question from a list of Documents."""
# Get the answer
chain = load_qa_with_sources_chain(
OpenAI(temperature=0,
openai_api_key=st.session_state.get("OPENAI_API_KEY")),
chain_type="stuff",
prompt=STUFF_PROMPT
)
# also returnig the text of the source used to form the answer
answer = chain(
{"input_documents": _docs, "question": query}
)
return answer
@st.cache_data
def get_sources(answer: Dict[str, Any], docs: List[Document]) -> List[Document]:
"""Gets the source documents for an answer."""
# Get sources for the answer
source_keys = [s for s in answer["output_text"].split("SOURCES: ")[-1].split(", ")]
source_docs = []
for doc in docs:
if doc.metadata["source"] in source_keys:
source_docs.append(doc)
return source_docs
def wrap_text_in_html(text: str) -> str:
"""Wraps each text block separated by newlines in <p> tags"""
if isinstance(text, list):
# Add horizontal rules between pages
text = "\n<hr/>\n".join(text)
return "".join([f"<p>{line}</p>" for line in text.split("\n")]) |