Spaces:
Running
on
CPU Upgrade
Running
on
CPU Upgrade
File size: 3,461 Bytes
4495c4a |
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 |
import os
import argparse
from tempfile import NamedTemporaryFile
from langchain.chains import create_retrieval_chain
from langchain.chains.combine_documents import create_stuff_documents_chain
from langchain_core.prompts import ChatPromptTemplate
from langchain_openai import ChatOpenAI
from langchain_community.document_loaders import PyPDFLoader
from langchain_community.vectorstores import FAISS
from langchain_openai import OpenAIEmbeddings
from langchain_text_splitters import RecursiveCharacterTextSplitter
def process_pdf(api_key, pdf_path):
os.environ["OPENAI_API_KEY"] = api_key
questions_path = "./Prompts/summary_tool_questions.md"
prompt_path = "./Prompts/summary_tool_system_prompt.md"
with open(pdf_path, "rb") as file:
with NamedTemporaryFile(delete=False, suffix=".pdf") as temp_pdf:
temp_pdf.write(file.read())
temp_pdf_path = temp_pdf.name
loader = PyPDFLoader(temp_pdf_path)
docs = loader.load()
text_splitter = RecursiveCharacterTextSplitter(chunk_size=3000, chunk_overlap=500)
splits = text_splitter.split_documents(docs)
vectorstore = FAISS.from_documents(
documents=splits, embedding=OpenAIEmbeddings(model="text-embedding-3-large")
)
retriever = vectorstore.as_retriever(search_kwargs={"k": 10})
if os.path.exists(prompt_path):
with open(prompt_path, "r") as file:
system_prompt = file.read()
else:
raise FileNotFoundError(f"The specified file was not found: {prompt_path}")
prompt = ChatPromptTemplate.from_messages(
[
("system", system_prompt),
("human", "{input}"),
]
)
llm = ChatOpenAI(model="gpt-4o")
question_answer_chain = create_stuff_documents_chain(llm, prompt, document_variable_name="context")
rag_chain = create_retrieval_chain(retriever, question_answer_chain)
if os.path.exists(questions_path):
with open(questions_path, "r") as file:
questions = [line.strip() for line in file.readlines() if line.strip()]
else:
raise FileNotFoundError(f"The specified file was not found: {questions_path}")
qa_results = []
for question in questions:
result = rag_chain.invoke({"input": question})
answer = result["answer"]
qa_text = f"### Question: {question}\n**Answer:**\n{answer}\n"
qa_results.append(qa_text)
os.remove(temp_pdf_path)
return qa_results
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="Generate a summary for a single PDF.")
parser.add_argument("api_key", type=str, help="OpenAI API Key")
parser.add_argument("pdf_path", type=str, help="Path to the PDF file")
args = parser.parse_args()
try:
results = process_pdf(args.api_key, args.pdf_path)
markdown_text = "\n".join(results)
# Define and create the output directory if it doesn't exist
output_folder = "CAPS_Summaries"
os.makedirs(output_folder, exist_ok=True)
# Save the results to a Markdown file
base_name = os.path.splitext(os.path.basename(args.pdf_path))[0]
output_file_path = os.path.join(output_folder, f"{base_name}_Summary.md")
with open(output_file_path, "w") as output_file:
output_file.write(markdown_text)
print(f"Summary saved to {output_file_path}")
except Exception as e:
print(f"An error occurred: {e}")
|