Sverd/image_captioner

1import cohere2from annoy import AnnoyIndex3import numpy as np4import dotenv5import os6import pandas as pd7 8dotenv.load_dotenv()9 10model_name = "embed-english-v3.0"11api_key = os.environ['COHERE_API_KEY']12input_type_embed = "search_document"13 14# Set up the cohere client15co = cohere.Client(api_key)16 17# Get the dataset of topics18topics = pd.read_csv("aicovers_topics.csv")19 20# Get the embeddings21list_embeds = co.embed(texts=list(topics['topic_cleaned']), model=model_name, input_type=input_type_embed).embeddings22 23# Create the search index, pass the size of embedding24search_index = AnnoyIndex(np.array(list_embeds).shape[1], metric='angular')25 26# Add vectors to the search index27for i in range(len(list_embeds)):28    search_index.add_item(i, list_embeds[i])29search_index.build(10)  # 10 trees30search_index.save('test.ann')31 32 33def topic_from_caption(caption):34    """35        Returns a topic from an uploaded list that is semantically similar to the input caption.36 37        Args:38        - caption (str): The image caption generated by MS Azure.39 40        Returns:41        - str: The extracted topic based on the provided caption.42        """43    input_type_query = "search_query"44    caption_embed = co.embed(texts=[caption], model=model_name, input_type=input_type_query).embeddings  # embeds a caption45    topic_ids = search_index.get_nns_by_vector(caption_embed[0], n=1, include_distances=True)  # retrieves the nearest category46    topic = topics.iloc[topic_ids[0]]['topic_cleaned'].to_string(index=False, header=False)47    return topic48 49