Function generate

utils/src/experiment/baseline_docpres.py:164–203

(model: Literal["Qwen2.5", "gpt"])

def generate(model: Literal["Qwen2.5", "gpt"]):
    if model == "Qwen2.5":
        llms.language_model = llms.qwen2_5
    elif model == "gpt":
        llms.language_model = llms.gpt4o

    print("Generating slides on baseline with ", llms.language_model.model)
    llm_name = llms.get_simple_modelname(llms.language_model)
    model = CLIPModel.from_pretrained("openai/clip-vit-large-patch14").to("cuda").eval()
    processor = CLIPProcessor.from_pretrained("openai/clip-vit-large-patch14")
    folders = list(glob("data/*/pdf/*"))
    progress = tqdm(total=len(folders))

    def process_folder(pdf_folder, model, processor):
        source_text = open(f"{pdf_folder}/source.md").read()
        bird_eye = json.load(open(f"{pdf_folder}/refined_doc.json"))
        images = json.load(open(f"{pdf_folder}/image_caption.json")).keys()
        output_dir = f"{pdf_folder}/docpres/{llm_name}"
        if os.path.exists(output_dir + "/final.jsonl"):
            progress.write(f"Skipping {pdf_folder}")
            progress.update(1)
            return
        try:
            generate_slides(
                output_dir,
                source_text,
                bird_eye,
                list(images),
                model,
                processor,
            )
            progress.update(1)
        except Exception as e:
            print(f"Error in {pdf_folder}: {e}")

    # for folder in folders:
    #     process_folder(folder, model, processor)

    with ThreadPoolExecutor() as executor:
        list(executor.map(lambda f: process_folder(f, model, processor), folders))


if __name__ == "__main__":