-
Notifications
You must be signed in to change notification settings - Fork 13
/
Copy pathindexer.py
316 lines (274 loc) · 13.7 KB
/
indexer.py
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
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
"""
Description: The script creates a faiss index from the analyze results json created with Azure AI Document Intelligence -> file variable.
The index files are written to faiss/{{sourcename}}. The extraced tables are added to the Documents according to the corresponding pages
"""
import json
import os
from langchain.docstore.document import Document
from langchain.embeddings.openai import OpenAIEmbeddings
from langchain.vectorstores import FAISS
import tablehelper as tb
def loadanalyzerjson(jsonfilepath):
import json
with open(jsonfilepath, encoding='utf-8') as json_file:
data = json.load(json_file)
if "analyzeResult" in data:
data=data["analyzeResult"]
print(jsonfilepath+" loaded with "+str(len(data['pages']))+" pages, "+str(len(data['paragraphs']))+" paragraphs and "+str(len(data['tables']))+" tables")
string=json.dumps(data)
string=string.replace("boundingRegions","bounding_regions")
string=string.replace("pageNumber","page_number")
string=string.replace("columnCount","column_count")
string=string.replace("rowCount","row_count")
string=string.replace("rowIndex","row_index")
string=string.replace("columnIndex","column_index")
string=string.replace("keyValuePairs","key_value_pairs")
#writing Text to file
txtfile=jsonfilepath+'.txt'
print("Writing plain text to "+txtfile)
with open(txtfile, "w", encoding='utf-8') as text_file:
text_file.write(data['content'])
return json.loads(string)
def gettokens(text):
import tiktoken
enc = tiktoken.encoding_for_model('gpt-3.5-turbo')
return len(enc.encode(text))
def createdocs(jsondata,maxtokensize,sourcename):
import tablehelper as tb
from langchain.docstore.document import Document
data=jsondata
spans={}
#collect spans from tables
for idx,tab in enumerate(data['tables']):
if(len(tab['spans'])==1):
key=tab['spans'][0]['offset']
spans[str(key)]=idx
else:
smallesoffset=9999999999999999999999999
for sp in tab['spans']:
if sp['offset']<smallesoffset:
smallesoffset=sp['offset']
spans[str(smallesoffset)]=idx
docs = []
pagenr=[]
sectionHeading=""
mdtextlist=[]
fullmdtext=""
endoftable=0
tablesearchkey=-1
largestdoc=0
pagecontent={}
mdtext=""
pagesources=[]
for i in range(1,len(data['pages'])+1):
pagecontent[str(i)]=""
print("Creating docs for "+sourcename+" with "+str(len(data['paragraphs']))+" paragraphes and "+str(len(data['tables']))+" tables")
#iterate over all paragraphes and create docs with mdtext seperated by sectionHeadings
for idx,paragraphes in enumerate(data['paragraphs']):
#when content is 7 Price conditions, then print content
if paragraphes['spans'][0]['offset']>=endoftable:
if 'role' in paragraphes:
#if paragraphes['role']=='sectionHeading' check if sectionHeading is the same as before:
if paragraphes['role']=='sectionHeading':
if sectionHeading!=paragraphes['content']:
#only create new doc if sectionHeading is not empty
if sectionHeading!='':
#build mdtext and create docs with content smaller than maxtokensize
mdtext="## "+sectionHeading+"\n\n"
#add content to pagecontent object
key=str(paragraphes['bounding_regions'][0]['page_number'])
if key in pagecontent:
pagecontent[key]=pagecontent[key]+"## "+paragraphes['content']+"\n\n"
else:
pagecontent[key]="## "+paragraphes['content']+"\n\n"
pagesources = []
for pid,md in enumerate(mdtextlist):
if gettokens(mdtext+md)<=maxtokensize:
mdtext=mdtext+md
if pagenr[pid] not in pagesources:
pagesources.append(pagenr[pid])
else:
if (gettokens(md)>maxtokensize):
tokens=gettokens(md)
if tokens>largestdoc:
largestdoc=tokens
docs.append(Document(page_content=md, metadata={"source": sourcename, "pages":[pagenr[pid]], "tokens":tokens}))
fullmdtext=fullmdtext+md
else:
tokens=gettokens(mdtext)
if tokens>largestdoc:
largestdoc=tokens
docs.append(Document(page_content=mdtext, metadata={"source": sourcename, "pages":pagesources, "tokens":tokens}))
#add to fullmdtext
fullmdtext=fullmdtext+mdtext
mdtext=md
pagesources = [pagenr[pid]]
#add last doc
if len(pagesources)>0:
fullmdtext=fullmdtext+mdtext
tokens=gettokens(mdtext)
if tokens>largestdoc:
largestdoc=tokens
docs.append(Document(page_content=mdtext, metadata={"source": sourcename, "pages":pagesources, "tokens":tokens}))
#reset mdtext and pagenr
mdtextlist=[]
pagenr=[]
#set new sectionHeading
sectionHeading=paragraphes['content']
else:
#add paragraphes to mdtext
mdtextlist.append(paragraphes['content']+"\n\n")
page=paragraphes['bounding_regions'][0]['page_number']
pagenr.append(page)
#add content to pagecontent object
key=str(paragraphes['bounding_regions'][0]['page_number'])
if key in pagecontent:
pagecontent[key]=pagecontent[key]+paragraphes['content']+"\n\n"
else:
pagecontent[key]=paragraphes['content']+"\n\n"
else:
mdtextlist.append(paragraphes['content']+"\n\n")
page=paragraphes['bounding_regions'][0]['page_number']
pagenr.append(page)
#add content to pagecontent object
key=str(paragraphes['bounding_regions'][0]['page_number'])
if key in pagecontent:
pagecontent[key]=pagecontent[key]+paragraphes['content']+"\n\n"
else:
pagecontent[key]=paragraphes['content']+"\n\n"
#add pagenr if not already in list
page=paragraphes['bounding_regions'][0]['page_number']
pagenr.append(page)
#add subsequent table if exists
searchkey=str(paragraphes['spans'][0]['offset']+paragraphes['spans'][0]['length']+1)
if tablesearchkey in spans or searchkey in spans:
i=spans[searchkey]
mdtextlist.append("\n\n"+tb.tabletomd(data['tables'][i])+"\n\n")
#add content to pagecontent object
key=str(paragraphes['bounding_regions'][0]['page_number'])
if key in pagecontent:
pagecontent[key]=pagecontent[key]+"\n\n"+tb.tabletomd(data['tables'][i])+"\n\n"
else:
pagecontent[key]="\n\n"+tb.tabletomd(data['tables'][i])+"\n\n"
if len(data['tables'][i]['spans'])>1:
smallesoffset=9999999999999999999999999
totallength=0
for sp in data['tables'][i]['spans']:
totallength=totallength+sp['length']
if sp['offset']<smallesoffset:
key=sp['offset']
smallesoffset=sp['offset']
endoftable=smallesoffset+totallength+1
tablesearchkey=smallesoffset+totallength+1
else:
endoftable=data['tables'][i]['spans'][0]['offset']+data['tables'][i]['spans'][0]['length']+1
tablesearchkey=data['tables'][i]['spans'][0]['offset']+data['tables'][i]['spans'][0]['length']+1
page=data['tables'][i]['bounding_regions'][0]['page_number']
pagenr.append(page)
for pid,md in enumerate(mdtextlist):
if gettokens(mdtext+md)<=maxtokensize:
mdtext=mdtext+md
if pagenr[pid] not in pagesources:
pagesources.append(pagenr[pid])
else:
if (gettokens(md)>maxtokensize):
tokens=gettokens(md)
if tokens>largestdoc:
largestdoc=tokens
docs.append(Document(page_content=md, metadata={"source": sourcename, "pages":[pagenr[pid]], "tokens":tokens}))
fullmdtext=fullmdtext+md
else:
tokens=gettokens(mdtext)
if tokens>largestdoc:
largestdoc=tokens
docs.append(Document(page_content=mdtext, metadata={"source": sourcename, "pages":pagesources, "tokens":tokens}))
#add to fullmdtext
fullmdtext=fullmdtext+mdtext
mdtext=md
pagesources = [pagenr[pid]]
#add last doc
if len(pagesources)>0:
#add to fullmdtext
fullmdtext=fullmdtext+mdtext
docs.append(Document(page_content=mdtext, metadata={"source": sourcename, "pages":pagesources, "tokens":gettokens(mdtext)}))
print("Created "+str(len(docs))+" docs with a total of "+str(gettokens(fullmdtext))+" tokens. Largest doc has "+str(largestdoc)+" tokens.")
return docs, pagecontent,fullmdtext
def create_embeddings(projectname,chunks,sourcename):
embeddings = OpenAIEmbeddings(deployment="text-embedding-ada-002", chunk_size=16)
vector_store = FAISS.from_documents(chunks, embeddings)
vector_store.save_local("projects/"+projectname+'/faiss/'+sourcename+'.json')
#vector_store = Chroma.from_documents(chunks, embeddings)
return vector_store
def createindex(projectname,jsonfile,sourcename,maxtokensize):
mdfile=jsonfile+'.md'
tablemdfile=jsonfile+".tables.md"
txtfile=jsonfile+'.txt'
contentjsonfile=jsonfile+".pagecontent.json"
keyvaluesjsonfile=jsonfile+".keyvalues.json"
#delete tablemd file if exists
if os.path.exists(tablemdfile):
os.remove(tablemdfile)
#delete md file if exists
if os.path.exists(mdfile):
os.remove(mdfile)
#delete txt file if exists
if os.path.exists(txtfile):
os.remove(txtfile)
#delete content.json file if exists
if os.path.exists(contentjsonfile):
os.remove(contentjsonfile)
#delete keyvalues.json file if exists
if os.path.exists(keyvaluesjsonfile):
os.remove(keyvaluesjsonfile)
data=loadanalyzerjson(jsonfile)
docs, pagecontent,fullmdtext=createdocs(data,maxtokensize,sourcename)
#writing fullmd to file
print("Writing markdown to "+mdfile)
with open(mdfile, "w", encoding='utf-8') as text_file:
text_file.write(fullmdtext)
#writing Text to file
print("Writing plain text to "+txtfile)
with open(txtfile, "w", encoding='utf-8') as text_file:
text_file.write(data['content'])
#writing pagecontent to file
with open(contentjsonfile, 'w', encoding='utf-8') as outfile:
json.dump(pagecontent, outfile, indent=4)
print(contentjsonfile+" created")
mdtext=""
for tabid, tab in enumerate(data['tables']):
mdtext=mdtext+"## Table "+str(tabid)+" from page "+str(tab['bounding_regions'][0]['page_number'])+"\n"+tb.tabletomd(tab)+"\n"
keyvalues={}
for i in range(1,len(data['pages'])+1):
keyvalues[str(i)]={}
#create keyvalues json
for keyvalue in data['key_value_pairs']:
pagekey=str(keyvalue['key']['bounding_regions'][0]['page_number'])
if 'value' in keyvalue:
if keyvalue['value'] is not None:
keyvalues[pagekey][keyvalue['key']['content']]=keyvalue['value']['content']
#writing keyvalues to file
with open(keyvaluesjsonfile, 'w', encoding='utf-8') as outfile:
json.dump(keyvalues, outfile, indent=4)
#writing to file
with open(tablemdfile, "w", encoding='utf-8') as text_file:
text_file.write(mdtext)
print(tablemdfile+" created")
print("Total Nr of Docs: "+str(len(docs)))
print("Creating FAISS Vector Store")
create_embeddings(projectname,docs,sourcename)
print("Done")
return sourcename+'.json'
if __name__ == "__main__":
import os
from dotenv import load_dotenv, find_dotenv
load_dotenv(find_dotenv(), override=True)
os.environ["OPENAI_API_BASE"] = os.environ["AZURE_OPENAI_ENDPOINT"]
os.environ["OPENAI_API_KEY"] = os.environ["AZURE_OPENAI_API_KEY"]
os.environ["OPENAI_API_VERSION"] = os.environ["AZURE_OPENAI_API_VERSION"]
os.environ["OPENAI_API_TYPE"] = "azure"
source="DWS Annual Report 2022_EN.pdf"
#load json file
file="projects/dws/files/source.json"
print("Source: "+source)
print("Loading JSON File")
createindex(file,source,debug=False,minchar=10,calculate_tokens=True)