Pash1986's picture
Upload app.py
5408474 verified
raw
history blame
2.02 kB
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()