Spaces:
Sleeping
Sleeping
import time | |
import gradio as gr | |
import os | |
import asyncio | |
from pymongo import MongoClient | |
from langchain_community.vectorstores import MongoDBAtlasVectorSearch | |
from langchain_openai 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 = 'embedded_movies' | |
collection = client[db_name][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='vector_index', text_key='plot', embedding_key='plot_embedding') | |
def get_movies(message, history): | |
# Use AsyncIO to run the similarity search in the background | |
# movies = vector_store.similarity_search(message, 3) | |
print ('Searching for: ' + message) | |
movies = vector_store.similarity_search(message, 3) | |
retrun_text = '' | |
for movie in movies: | |
retrun_text = retrun_text + 'Title : ' + movie.metadata['title'] + '\n------------\n' + 'Plot: ' + movie.page_content + '\n\n' | |
for i in range(len(retrun_text)): | |
time.sleep(0.05) | |
yield "Found: " + "\n\n" + retrun_text[: 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() |