FraleyLabAttachmentBot / ObtainDataEmbedding.py
AjithKSenthil's picture
modified it to use our data now
8b16ee5
# ObtainDataEmbedding.py
# imports
import openai
import pandas as pd
import tiktoken
from openai.embeddings_utils import get_embedding
import config
# set your API key
openai.api_key = "your openai api key"
# embedding model parameters
embedding_model = "text-embedding-ada-002"
embedding_encoding = "cl100k_base" # this the encoding for text-embedding-ada-002
max_tokens = 8000 # the maximum for text-embedding-ada-002 is 8191
# load & inspect dataset
input_datapath = "data/attachmentchatdata-formated.csv"
df = pd.read_csv(input_datapath, index_col=0)
df = df[["userid", "chathistory", "avoide", "avoida", "avoidb", "avoidc", "avoidd", "anxietye", "anxietya", "anxietyb", "anxietyc", "anxietyd"]]
df = df.dropna()
df.head(2)
# Filter out chat transcripts that are too long to embed, estimate for the maximum number of words would be around 1638 words (8191 tokens / 5).
encoding = tiktoken.get_encoding(embedding_encoding)
df["n_tokens"] = df.chathistory.apply(lambda x: len(encoding.encode(x)))
df = df[df.n_tokens <= max_tokens]
len(df)
# Ensure you have your API key set in your environment per the README: https://github.com/openai/openai-python#usage
# This may take a few minutes
df["embedding"] = df.chathistory.apply(lambda x: get_embedding(x, engine=embedding_model))
df.to_csv("data/chat_transcripts_with_embeddings_and_scores.csv")
# Please replace "data/chat_transcripts.csv" with the path to your actual data file. Also, replace 'ChatTranscript', 'Attachment', 'Avoidance' with the actual column names of your chat transcripts and attachment scores in your data file.
# Also, remember to set the API key for OpenAI in your environment before running the get_embedding function.