Large Scale Embedding Generation with Qwen3 0.6B

# Authenticate
credentials = service_account.Credentials.from_service_account_info(
    json.loads(os.environ['SERVICE_ACCOUNT_JSON'])
)
project_id = 'xxxxx'
client = bigquery.Client(project=project_id, credentials=credentials)

# Run query
query_job = client.query(query_full_text)
results = query_job.result()
patents_df = pl.from_arrow(results.to_arrow())

# Process patents for embedding
all_documents = []
for row in patents_df.iter_rows(named=True):
    all_documents.extend(prepare_patent_for_embedding(row))

docs_df = pl.DataFrame({
    'doc_id': [doc['doc_id'] for doc in all_documents],
    'patent_id': [doc['patent_id'] for doc in all_documents],
    'section': [doc['section'] for doc in all_documents],
    'text': [doc['text'] for doc in all_documents],
    'metadata': [json.dumps(doc['metadata']) for doc in all_documents]
})

# Create output directory
output_dir = "apple_patents_documents_parquet"
os.makedirs(output_dir, exist_ok=True)

# Calculate rows per chunk for ~500MB files
total_size_mb = docs_df.estimated_size() / (1024 * 1024)
num_files = max(1, int(total_size_mb / 500) + 1)
rows_per_chunk = len(docs_df) // num_files + 1

# Split and save
for i in range(0, len(docs_df), rows_per_chunk):
    chunk = docs_df[i:i+rows_per_chunk]
    filename = f"{output_dir}/documents_part_{i//rows_per_chunk:03d}.parquet"
    chunk.write_parquet(filename, compression='snappy')

print(f"Saved {len(docs_df)} documents across {num_files} files to {output_dir}/")

dataset_id = sutro.create_dataset()
sutro.upload_to_dataset(dataset_id, output_dir)

job_id = sutro.infer(
    dataset_id,
    column="text",
    model="qwen-3-embedding-0.6b",
    job_priority=1,
)

sutro.await_job_completion(job_id, obtain_results=False)

results_dir='/mnt/cooper-notebook-volume/apple_patents_embeddings'
os.makedirs(results_dir, exist_ok=True)
sutro.download_from_dataset(dataset_id, output_path=results_dir)
job_results_df = pl.read_parquet(results_dir)

print(f"Total patents: {len(job_results_df)}")
print(f"Columns: {job_results_df.columns}")
# Total patents: 4039988
# Columns: ['text', 'job-76844041-b2bf-4248-9603-b7f750231b34']

vector_col_name = 'job-76844041-b2bf-4248-9603-b7f750231b34'

# Combine the embeddings with docs_df containing attributing IDs and metadata
combined_df = docs_df.with_columns(
    job_results_df[vector_col_name]
)

vector_col_name = 'job-76844041-b2bf-4248-9603-b7f750231b34'
# We're using an im-memory DB here for convenience, but Qdrant Cloud
# can be faster to upload to and will persist the embeddings as well
client = QdrantClient(":memory:")
collection_name = "apple_patents_collection"

# Create the collection, inferring vector size from the first
# row of the DataFrame
client.recreate_collection(
    collection_name=collection_name,
    vectors_config=models.VectorParams(
        size=len(combined_df[vector_col_name][0]),
        distance=models.Distance.COSINE
    ),
)

BATCH_SIZE = 8192

# Loop through the DataFrame in batches
for i in tqdm(range(0, len(combined_df), BATCH_SIZE), desc="Uploading to Qdrant"):
    batch_df = combined_df.slice(i, BATCH_SIZE)

    points = [
        models.PointStruct(
            id=row['doc_id'],
            vector=row[vector_col_name],
            payload={
                'patent_id': row['patent_id'],
                'section': row['section'],
                'text': row['text'],
                **json.loads(row['metadata'])
            }
        ) for row in batch_df.to_dicts()
    ]

    client.upload_points(
        collection_name=collection_name,
        points=points,
        wait=True,
        parallel=4
    )

print(f"Finished uploading all {len(combined_df)} points.")

def search_patents(query_text, top_k=5):
    # Pick your real time provider of choice, we have heard great
    # things (latency & consistency wise) about Vertex, but Baseten
    # is the easier choice
    # https://www.baseten.co/library/qwen3-06b-embedding/
    query_embedding = real_time_api(
        text=query_text,
        model="qwen-3-embedding-0.6b"
    )

    # Search in Qdrant
    results = client.search(
        collection_name=collection_name,
        query_vector=query_embedding,
        limit=top_k
    )

    for result in results:
        print(f"Patent: {result.payload['patent_id']}")
        print(f"Section: {result.payload['section']}")
        print(f"Score: {result.score:.3f}")
        print(f"Text: {result.payload['text'][:200]}...")
        print("-" * 80)

# Let's test it out!
search_patents("wireless charging efficiency improvements")