(pdf_folder: str, rank: int)
| 86 | |
| 87 | |
| 88 | def prepare_pdf_folder(pdf_folder: str, rank: int): |
| 89 | image_model = get_image_model(f"cuda:{rank % device_count}") |
| 90 | if not pexists(pjoin(pdf_folder, "source.md")): |
| 91 | return |
| 92 | if not pexists(pjoin(pdf_folder, "image_caption.json")): |
| 93 | images_embeddings = get_image_embedding(pdf_folder, *image_model) |
| 94 | images = [pjoin(pdf_folder, image) for image in images_embeddings] |
| 95 | if len(images_embeddings) == 0: |
| 96 | rm_folder(pdf_folder) |
| 97 | return |
| 98 | similarity_matrix = images_cosine_similarity(list(images_embeddings.values())) |
| 99 | for i in range(len(similarity_matrix)): |
| 100 | for j in range(i + 1, len(similarity_matrix)): |
| 101 | if similarity_matrix[i][j] > 0.85: |
| 102 | if pexists(images[i]): |
| 103 | os.remove(images[i]) |
| 104 | break |
| 105 | images = [image for image in images if pexists(image)] |
| 106 | image_stats = {} |
| 107 | caption_prompt = open("prompts/caption.txt").read() |
| 108 | for image in images: |
| 109 | image_stats[image] = llms.vision_model(caption_prompt, image) |
| 110 | print(image_stats[image]) |
| 111 | with open(pjoin(pdf_folder, "image_caption.json"), mode="w") as f: |
| 112 | json.dump(image_stats, f, indent=4, ensure_ascii=False) |
| 113 | |
| 114 | if not pexists(pjoin(pdf_folder, "refined_doc.json")): |
| 115 | text_content = open(pjoin(pdf_folder, "source.md")).read() |
| 116 | text_content = markdown_clean_pattern.sub("", text_content) |
| 117 | template = Template(open("prompts/document_refine.txt").read()) |
| 118 | doc_json = llms.language_model( |
| 119 | template.render(markdown_document=text_content), return_json=True |
| 120 | ) |
| 121 | json.dump( |
| 122 | doc_json, |
| 123 | open(pjoin(pdf_folder, "refined_doc.json"), "w"), |
| 124 | indent=4, |
| 125 | ensure_ascii=False, |
| 126 | ) |
| 127 | |
| 128 | |
| 129 | def filter_slide(slide: SlidePage): |
nothing calls this directly
no test coverage detected