Spaces:
Sleeping
Sleeping
| 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() |