Spaces:
Sleeping
Sleeping
File size: 2,711 Bytes
1ec4e65 874c091 1ec4e65 0e51c38 1ec4e65 |
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 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 |
import os
import time
import pickle
import streamlit as st
from dotenv import load_dotenv
from pinecone import Pinecone, ServerlessSpec
from utils import load_pickle, initialize_embedding_model
from langchain_community.retrievers import BM25Retriever
from langchain_pinecone import PineconeVectorStore
from langchain.retrievers import EnsembleRetriever
from langchain.tools.retriever import create_retriever_tool
# Load .env file
load_dotenv()
# Constants
INDEX_NAME = "veda-index-v2"
MODEL_NAME = "BAAI/bge-large-en-v1.5"
# Initialize Pinecone client
os.environ['PINECONE_API_KEY'] = os.getenv("PINECONE_API_KEY")
PINECONE_API_KEY = os.getenv("PINECONE_API_KEY")
pc = Pinecone(api_key=PINECONE_API_KEY)
#@st.cache_resource
def create_or_load_index():
# Check if index already exists
if INDEX_NAME not in pc.list_indexes().names():
# Create index if it does not exist
pc.create_index(
INDEX_NAME,
dimension=1024,
metric='dotproduct',
spec=ServerlessSpec(
cloud="aws",
region="us-east-1"
)
)
# Wait for index to be initialized
while not pc.describe_index(INDEX_NAME).status['ready']:
time.sleep(1)
# Connect to index
return pc.Index(INDEX_NAME)
# Load documents
docs = load_pickle("ramana_docs_ids.pkl")
# Initialize embedding model
embedding = initialize_embedding_model(MODEL_NAME)
# Create or load index
index = create_or_load_index()
# Initialize BM25 retriever
bm25_retriever = BM25Retriever.from_texts(
[text['document'].page_content for text in docs],
metadatas=[text['document'].metadata for text in docs]
)
bm25_retriever.k = 2
# Switch back to normal index for LangChain
vector_store = PineconeVectorStore(index, embedding)
retriever = vector_store.as_retriever(search_type="mmr")
# Initialize the ensemble retriever
ensemble_retriever = EnsembleRetriever(
retrievers=[bm25_retriever, retriever], weights=[0.2, 0.8]
)
vector_tools = create_retriever_tool(
retriever = ensemble_retriever,
name = "vector_retrieve",
description="Search and return documents related user query from the vector index.",
)
from langchain import hub
prompt = hub.pull("hwchase17/openai-tools-agent")
prompt.messages
from langchain.agents import AgentExecutor, create_openai_tools_agent
from langchain_openai import ChatOpenAI
import streamlit as st
os.environ["OPENAI_API_KEY"] = st.secrets["OPENAI_API_KEY"]
#load llm model
llm_AI4 = ChatOpenAI(model="gpt-4-1106-preview", temperature=0)
agent = create_openai_tools_agent(llm_AI4, [vector_tools], prompt)
agent_executor = AgentExecutor(agent=agent, tools=[vector_tools])
|