File size: 457 Bytes
c254fea
90aed10
c254fea
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
import pandas as pd
from utils.openai_api import get_embedding

def compute_doc_embeddings(df: pd.DataFrame) -> dict[tuple[str, str], list[float]]:
    """
    Create an embedding for each row in the dataframe using the OpenAI Embeddings API.
    
    Return a dictionary that maps between each embedding vector and the index of the row that it corresponds to.
    """
    return {
        idx: get_embedding(r.description) for idx, r in df.iterrows()
    }