Usando os embeddings mais recentes da OpenAI em um sistema RAG com o MongoDB

1
!pip install datasets pandas openai pymongo 

1
from datasets import load_dataset
2
import pandas as pd
3
4
# <https://huggingface.co/datasets/AIatMongoDB/embedded_movies>
5
dataset = load_dataset("AIatMongoDB/embedded_movies")
6
7
# Convert the dataset to a pandas dataframe
8
dataset_df = pd.DataFrame(dataset['train'])
9
10
dataset_df.head(5)

1
# Remove data point where plot column is missing
2
dataset_df = dataset_df.dropna(subset=['plot'])
3
print("\\nNumber of missing values in each column after removal:")
4
print(dataset_df.isnull().sum())
5
6
# Remove the plot_embedding from each data point in the dataset as we are going to create new embeddings with the new OpenAI embedding Model "text-embedding-3-small"
7
dataset_df = dataset_df.drop(columns=['plot_embedding'])
8
dataset_df.head(5)

1
python
2
import openai
3
from google.colab import userdata
4
5
openai.api_key = userdata.get("open_ai")
6
7
EMBEDDING_MODEL = "text-embedding-3-small"
8
9
def get_embedding(text):
10
    """Generate an embedding for the given text using OpenAI's API."""
11
12
    # Check for valid input
13
    if not text or not isinstance(text, str):
14
        return None
15
16
    try:
17
        # Call OpenAI API to get the embedding
18
        embedding = openai.embeddings.create(input=text, model=EMBEDDING_MODEL).data[0].embedding
19
        return embedding
20
    except Exception as e:
21
        print(f"Error in get_embedding: {e}")
22
        return None
23
24
dataset_df["plot_embedding_optimised"] = dataset_df['plot'].apply(get_embedding)
25
26
dataset_df.head()

1
import pymongo
2
from google.colab import userdata
3
4
def get_mongo_client(mongo_uri):
5
  """Establish connection to the MongoDB."""
6
  try:
7
    client = pymongo.MongoClient(mongo_uri)
8
    print("Connection to MongoDB successful")
9
    return client
10
  except pymongo.errors.ConnectionFailure as e:
11
    print(f"Connection failed: {e}")
12
    return None
13
14
mongo_uri = userdata.get('MONGO_URI')
15
if not mongo_uri:
16
  print("MONGO_URI not set in environment variables")
17
18
mongo_client = get_mongo_client(mongo_uri)
19
20
# Ingest data into MongoDB
21
db = mongo_client['movies']
22
collection = db['movie_collection'] 
23
24
documents = dataset_df.to_dict('records')
25
collection.insert_many(documents)
26
27
print("Data ingestion into MongoDB completed")

1
{
2
 "fields": [
3
   {
4
     "numDimensions": 1536,
5
     "path": "plot_embedding_optimised",
6
     "similarity": "cosine",
7
     "type": "vector"
8
   }
9
 ]
10
}

1
def vector_search(user_query, collection):
2
    """
3
    Perform a vector search in the MongoDB collection based on the user query.
4
5
    Args:
6
    user_query (str): The user's query string.
7
    collection (MongoCollection): The MongoDB collection to search.
8
9
    Returns:
10
    list: A list of matching documents.
11
    """
12
13
    # Generate embedding for the user query
14
    query_embedding = get_embedding(user_query)
15
16
    if query_embedding is None:
17
        return "Invalid query or embedding generation failed."
18
19
    # Define the vector search pipeline
20
    pipeline = [
21
        {
22
            "$vectorSearch": {
23
                "index": "vector_index",
24
                "queryVector": query_embedding,
25
                "path": "plot_embedding_optimised",
26
                "numCandidates": 150,  # Number of candidate matches to consider
27
                "limit": 5  # Return top 5 matches
28
            }
29
        },
30
        {
31
            "$project": {
32
                "_id": 0,  # Exclude the _id field
33
                "plot_embedding_opitimzed": 0,  # Exclude the plot_embedding_opitimzed field
34
                "plot": 1,  # Include the plot field
35
                "title": 1,  # Include the title field
36
                "genres": 1, # Include the genres field
37
                "score": {
38
                    "$meta": "vectorSearchScore"  # Include the search score
39
                }
40
            }
41
        }
42
    ]
43
44
    # Execute the search
45
    results = collection.aggregate(pipeline)
46
    return list(results)

1
def handle_user_query(query, collection):
2
3
  get_knowledge = vector_search(query, collection)
4
5
  search_result = ''
6
  for result in get_knowledge:
7
      search_result += f"Title: {result.get('title', 'N/A')}, Plot: {result.get('plot', 'N/A')}\\n"
8
9
  completion = openai.chat.completions.create(
10
      model="gpt-3.5-turbo",
11
      messages=[
12
          {"role": "system", "content": "You are a movie recommendation system."},
13
          {"role": "user", "content": "Answer this user query: " + query + " with the following context: " + search_result}
14
      ]
15
  )
16
17
  return (completion.choices[0].message.content), search_result
18
19
Conduct query with retrieval of sources
20
query = "What is the best romantic movie to watch?"
21
response, source_information = handle_user_query(query, collection)
22
23
print(f"Response: {response}")
24
print(f"Source Information: \\n{source_information}")

1
Response: Based on the context provided, the best romantic movie to watch would be "Gorgeous". It revolves around a romantic girl who travels to Hong Kong in search of true love but unexpectedly falls for a kind-hearted professional fighter.
2
Source Information: 
3
Title: Run, Plot: This action movie is filled with romance and adventure. As Abhisek fights for his life against the forces of crime and injustice, he meets Bhoomika, who captures his heart.
4
Title: China Girl, Plot: A modern day Romeo & Juliet story is told in New York when an Italian boy and a Chinese girl become lovers, causing a tragic conflict between ethnic gangs.
5
Title: Gorgeous, Plot: A romantic girl travels to Hong Kong in search of certain love but instead meets a kind-hearted professional fighter with whom she begins to fall for instead.
6
Title: Once a Thief, Plot: A romantic and action packed story of three best friends, a group of high end art thieves, who come into trouble when a love-triangle forms between them.
7
Title: House of Flying Daggers, Plot: A romantic police captain breaks a beautiful member of a rebel group out of prison to help her rejoin her fellows, but things are not what they seem.