Spaces:
Running
Running
import time | |
import gradio as gr | |
import os | |
import asyncio | |
from pymongo import MongoClient | |
from langchain.vectorstores import MongoDBAtlasVectorSearch | |
from langchain.embeddings import OpenAIEmbeddings | |
from langchain.llms import OpenAI | |
from langchain.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() |