Generating data for embedding training materials¶

def download_pdfs_from_zenodo(record_id):
    """Download PDFs from Zenodo record."""
    base_url = f"https://zenodo.org/api/records/{record_id}"
    response = requests.get(base_url)
    data = response.json()
    
    if not os.path.exists('downloads'):
        os.makedirs('downloads')
    
    files_info = []
    for file in data['files']:
        if file['key'].endswith('.pdf'):
            download_url = file['links']['self']
            filename = record_id + "_" + file['key']
            filepath = os.path.join('downloads', filename)

            if not os.path.exists(filepath):
                # Download file
                response = requests.get(download_url)
                with open(filepath, 'wb') as f:
                    f.write(response.content)
            
            files_info.append({'filename': filename, 'url': download_url})
    
    return files_info


# Download PDFs
files_info = download_pdfs_from_zenodo('12623730')
files_info

def resize_image(image, height):
    """
    Resize the image to the specified height while maintaining aspect ratio.

    Parameters
    ----------
    image : PIL.Image.Image
        The image to resize.
    height : int
        The desired height in pixels.
    
    Returns
    -------
    PIL.Image.Image
        The resized image.
    """
    aspect_ratio = image.width / image.height
    new_width = int(aspect_ratio * height)
    return image.resize((new_width, height), Image.LANCZOS)

def process_pdf(pdf_info):
    """Process PDF file to extract images and text."""
    filename = pdf_info['filename']
    filepath = os.path.join('downloads', filename)
    base_name = os.path.splitext(filename)[0]
    
    if not os.path.exists('downloads'):
        os.makedirs('downloads')
    if not os.path.exists('downloads'):
        os.makedirs('downloads')

    # Set your OpenAI API key
    client = OpenAI()
    
    # Convert PDF pages to images
    images = [resize_image(i, height=300) for i in convert_from_path(filepath)]
    
    # Extract text from PDF
    pdf_reader = PyPDF2.PdfReader(filepath)
    
    page_data = []
    
    for i, image in enumerate(images):
        # Save image
        png_filename = f"{base_name}_{i}.png"
        png_path = os.path.join('downloads', png_filename)
        image.save(png_path)
        
        # Save text
        txt_filename = f"{base_name}_{i}.txt"
        txt_path = os.path.join('downloads', txt_filename)
        text = pdf_reader.pages[i].extract_text()
        if not os.path.exists(txt_filename):
            with open(txt_path, 'w', encoding='utf-8') as f:
                f.write(text)
            
        # Get embedding
        response = client.embeddings.create(
            input=text,
            model="text-embedding-ada-002"
        )
        embedding_vector = response.data[0].embedding
        
        page_data.append({
            'filename': filename,
            'url': pdf_info['url'],
            'page_index': i,
            'text': text,
            'png_filename': png_filename,
            'txt_filename': txt_filename,
            'embedding_vector': embedding_vector
        })
     
    return page_data

# Process all PDFs
all_page_data = []
for pdf_info in files_info:
    all_page_data.extend(process_pdf(pdf_info))

# Create DataFrame
df = pd.DataFrame(all_page_data)
df

# Convert embedding vectors to numpy array for UMAP
embeddings = np.array(df['embedding_vector'].tolist())

# Apply UMAP
reducer = umap.UMAP(n_components=2, random_state=42)
umap_embeddings = reducer.fit_transform(embeddings)

df['UMAP0'] = umap_embeddings[:, 0]
df['UMAP1'] = umap_embeddings[:, 1]

df

import stackview
stackview.scatterplot(df, column_x="UMAP0", column_y="UMAP1")

df["selection"]

import yaml

# Convert DataFrame to dictionary
data_dict = df.to_dict()

# Save as YAML file
with open('data.yml', 'w') as file:
    yaml.dump(data_dict, file)

print("DataFrame saved as data.yml")

import yaml

# Convert DataFrame to dictionary
data_dict = df[["filename", "png_filename", "text", "url", "page_index", "UMAP0", "UMAP1"]].to_dict()

# Save as YAML file
with open('data_png_umap.yml', 'w') as file:
    yaml.dump(data_dict, file)

print("DataFrame saved as data.yml")