cckp_publication_abstracts.py

import synapseclient
import boto3
import json
import time
from concurrent.futures import ThreadPoolExecutor, as_completed
from hallmarks_llm import extract_cancer_hallmarks
from tqdm import tqdm

# Initialize Synapse and AWS
syn = synapseclient.Synapse()
syn.login()

session = boto3.Session(profile_name="htan-dev")
bedrock = session.client(service_name="bedrock-runtime")

# Fetch the table
table = syn.tableQuery("SELECT * FROM syn21868591")
df = table.asDataFrame()

# Initialize columns for hallmarks and scores
df["hallmarks"] = None
df["scores"] = None


def extract_hallmarks_with_retries(index, row, max_retries=5, base_wait_time=2):
    """
    Extract cancer hallmarks from the given abstract with retry logic for throttling errors.
    """
    retries = 0
    while retries < max_retries:
        try:
            # Extract hallmarks from the abstract
            result = extract_cancer_hallmarks(row["abstract"])
            hallmarks = [item["hallmark"] for item in result["extracted_hallmarks"]]
            scores = [
                item["confidence_score"] for item in result["extracted_hallmarks"]
            ]
            return index, json.dumps(hallmarks), json.dumps(scores)
        except Exception as e:
            if "ThrottlingException" in str(e):
                retries += 1
                wait_time = base_wait_time * (2 ** (retries - 1))  # Exponential backoff
                print(
                    f"Throttling error on row {index}. Retrying in {wait_time}s... "
                    f"(Attempt {retries}/{max_retries})"
                )
                time.sleep(wait_time)
            else:
                print(f"Unhandled error on row {index}: {e}")
                return index, None, None
    print(f"Max retries exceeded for row {index}")
    return index, None, None


# Process rows in parallel
with ThreadPoolExecutor(max_workers=30) as executor:
    futures = {
        executor.submit(extract_hallmarks_with_retries, i, row): i
        for i, row in df.iterrows()
    }

    # Display progress bar
    for future in tqdm(
        as_completed(futures), total=len(futures), desc="Processing abstracts"
    ):
        index = futures[future]
        try:
            index, hallmarks, scores = future.result()
            df.at[index, "hallmarks"] = hallmarks
            df.at[index, "scores"] = scores
        except Exception as e:
            print(f"Error processing row {index}: {e}")

# Save to CSV
df.to_csv("cckp_publication_abstracts.csv", index=False, quoting=1)

print("Processing complete. Data saved to 'cckp_publication_abstracts.csv'.")