Pash1986's picture
Upload 2 files
e2275e6 verified
raw
history blame
1.84 kB
import time
import gradio as gr
import os
import asyncio
from pymongo import MongoClient
from langchain_community.vectorstores import MongoDBAtlasVectorSearch
from langchain_community.embeddings import OpenAIEmbeddings
# from langchain_community.llms import OpenAI
# from langchain_community.prompts import PromptTemplate
# from langchain.chains import LLMChain
import json
## Connect to MongoDB Atlas local cluster
MONGODB_ATLAS_CLUSTER_URI = os.getenv('MONGODB_ATLAS_CLUSTER_URI')
client = MongoClient(MONGODB_ATLAS_CLUSTER_URI)
db_name = 'sample_mflix'
collection_name = 'movies'
collection = client[db_name][collection_name]
## Create a collection and insert data
print ('Creating collection movies')
client[db_name].create_collection(collection_name)
## Create a vector search index
print ('Creating vector search index')
collection.create_search_index(model={"definition": {"mappings":{
"dynamic":True,
"fields": {
"plot_embedding": {
"type": "knnVector",
"dimensions": 1536,
"similarity": "euclidean"
}
}
}}, "name":'default'})
# sleep for minute
print ('Waiting for vector index on field "embedding" to be created')
time.sleep(60)
vector_store = MongoDBAtlasVectorSearch(embedding=OpenAIEmbeddings(), collection=collection, index_name='default', text_key='plot', embedding_key='plot_embedding')
def get_movies(message, history):
movies = vector_store.similarity_search(message, 3)
for movie in movies:
for i in range(len(movie.metadata['title'])):
time.sleep(0.05)
yield "Movie " + i + " : Title - " + movie.metadata['title'][: i+1]
demo = gr.ChatInterface(get_movies, examples=["What movies are scary?", "Find me a comedy", "Movies for kids"], title="Movies Atlas Vector Search", submit_btn="Search").queue()
if __name__ == "__main__":
demo.launch()