Spaces:
Sleeping
Sleeping
File size: 2,018 Bytes
8f835ef 90f787c 8f835ef e2275e6 cacb7f3 47f7c8e e2275e6 8f835ef cacb7f3 8f835ef ac6abfb 8f835ef ac6abfb 8f835ef ac6abfb 8f835ef cacb7f3 8f835ef 5408474 cacb7f3 8f835ef |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 |
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() |