XicoC commited on
Commit
7193baa
1 Parent(s): 91402bc

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +24 -3
app.py CHANGED
@@ -14,6 +14,7 @@ from aimakerspace.openai_utils.embedding import EmbeddingModel
14
  from aimakerspace.vectordatabase import VectorDatabase
15
  from aimakerspace.openai_utils.chatmodel import ChatOpenAI
16
  import chainlit as cl
 
17
 
18
  system_template = """\
19
  Use the following context to answer a users question. If you cannot find the answer in the context, say you don't know the answer."""
@@ -76,6 +77,23 @@ def process_text_file(file: AskFileResponse):
76
  return texts
77
 
78
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
79
  @cl.on_chat_start
80
  async def on_chat_start():
81
  files = None
@@ -84,7 +102,7 @@ async def on_chat_start():
84
  while files == None:
85
  files = await cl.AskFileMessage(
86
  content="Please upload a Text File file to begin!",
87
- accept=["text/plain"],
88
  max_size_mb=2,
89
  timeout=180,
90
  ).send()
@@ -97,7 +115,10 @@ async def on_chat_start():
97
  await msg.send()
98
 
99
  # load the file
100
- texts = process_text_file(file)
 
 
 
101
 
102
  print(f"Processing {len(texts)} text chunks")
103
 
@@ -129,4 +150,4 @@ async def main(message):
129
  async for stream_resp in result["response"]:
130
  await msg.stream_token(stream_resp)
131
 
132
- await msg.send()
 
14
  from aimakerspace.vectordatabase import VectorDatabase
15
  from aimakerspace.openai_utils.chatmodel import ChatOpenAI
16
  import chainlit as cl
17
+ from langchain_community.document_loaders import PyPDFLoader
18
 
19
  system_template = """\
20
  Use the following context to answer a users question. If you cannot find the answer in the context, say you don't know the answer."""
 
77
  return texts
78
 
79
 
80
+ def process_pdf_file(file: AskFileResponse):
81
+ import tempfile
82
+
83
+ with tempfile.NamedTemporaryFile(
84
+ mode="w", delete=False, suffix=".pdf"
85
+ ) as temp_file:
86
+ temp_file_path = temp_file.name
87
+
88
+ with open(temp_file_path, "wb") as f:
89
+ f.write(file.content)
90
+
91
+ pdf_loader = PyPDFLoader(temp_file_path)
92
+ documents = pdf_loader.load()
93
+ texts = text_splitter.split_texts(documents)
94
+ return texts
95
+
96
+
97
  @cl.on_chat_start
98
  async def on_chat_start():
99
  files = None
 
102
  while files == None:
103
  files = await cl.AskFileMessage(
104
  content="Please upload a Text File file to begin!",
105
+ accept=["text/plain", "application/pdf"],
106
  max_size_mb=2,
107
  timeout=180,
108
  ).send()
 
115
  await msg.send()
116
 
117
  # load the file
118
+ if file.path.endswith(".pdf"):
119
+ texts = process_pdf_file(file)
120
+ else:
121
+ texts = process_text_file(file)
122
 
123
  print(f"Processing {len(texts)} text chunks")
124
 
 
150
  async for stream_resp in result["response"]:
151
  await msg.stream_token(stream_resp)
152
 
153
+ await msg.send()