Spaces:
Running
Running
import streamlit as st | |
import os | |
import time | |
import numpy as np | |
import pandas as pd | |
def add_custom_css(): | |
st.markdown(""" | |
<style> | |
.container { | |
text-align: center; | |
background-color: #f0f0f0; | |
padding: 20px; | |
} | |
.big-font { | |
font-size: 50px; | |
color: #4CAF50; | |
} | |
.progress-bar { | |
margin-top: 20px; | |
} | |
</style> | |
""", unsafe_allow_html=True) | |
if 'packages_installed' not in st.session_state: | |
st.info("Installing required packages...") | |
os.system("pip install -U sentence-transformers") | |
os.system("pip install pinecone-client") | |
st.session_state['packages_installed'] = True | |
from sentence_transformers import SentenceTransformer | |
from pinecone import Pinecone, ServerlessSpec, PodSpec | |
if 'pc' not in st.session_state: | |
use_serverless = False | |
# Configure Pinecone client | |
api_key = os.environ.get('PINECONE_API_KEY', '28b0fd5a-fdfb-422d-9a44-c0ec09a25074') | |
environment = os.environ.get('PINECONE_ENVIRONMENT', 'gcp-starter') | |
st.session_state['pc'] = Pinecone(api_key=api_key) | |
if use_serverless: | |
spec = ServerlessSpec(cloud='gcp', region='asia-southeast1-gcp') | |
else: | |
spec = PodSpec(environment=environment) | |
if 'model' not in st.session_state: | |
st.session_state['model'] = SentenceTransformer('intfloat/e5-small') | |
index_name = 'dataset' | |
if index_name not in st.session_state.pc.list_indexes().names(): | |
dimensions = 384 | |
st.session_state.pc.create_index( | |
name=index_name, | |
dimension=dimensions, | |
metric='cosine', | |
spec=spec | |
) | |
# Wait until index is ready | |
while not st.session_state.pc.describe_index(index_name).status['ready']: | |
time.sleep(1) | |
if 'index' not in st.session_state: | |
st.session_state['index'] = st.session_state.pc.Index(index_name) | |
# Function to process data and insert into Pinecone index | |
def process_data(data, namespace): | |
input_texts = data['Query'] | |
progress_bar = st.progress(0) | |
total_chunks = len(data) // 1000 + 1 | |
for chunk_start in range(0, len(data), 1000): | |
chunk_end = min(chunk_start + 1000, len(data)) | |
chunk = data.iloc[chunk_start:chunk_end] | |
# Generate embeddings for the current chunk | |
chunk_embeddings = [st.session_state.model.encode(query, normalize_embeddings=True) for query in chunk['Query']] | |
chunk['embedding'] = chunk_embeddings | |
# Upsert embeddings | |
st.session_state.index.upsert(vectors=zip(chunk['id'], chunk['embedding']), namespace=namespace) | |
# Update progress bar | |
progress = (chunk_end / len(data)) * 100 | |
progress_bar.progress(int(progress)) | |
def load_and_process_data(file): | |
data = pd.read_csv(file) | |
data['id'] = data.index.astype(str) | |
namespace = file.name[:15] # Use first 15 characters of file name as namespace | |
if 'embeddings_done' not in st.session_state: | |
process_data(data, namespace) | |
st.session_state['embeddings_done'] = True | |
return data, namespace | |
def main(): | |
add_custom_css() | |
st.markdown(""" | |
<div class='container'> | |
<h1 class='big-font'>Semantic Search Engine</h1> | |
</div> | |
""", unsafe_allow_html=True) | |
# Use session state to retain information across interactions | |
if 'namespace' not in st.session_state: | |
st.session_state.namespace = None | |
if 'df' not in st.session_state: | |
st.session_state.df = None | |
uploaded_file = st.file_uploader("Upload dataset (CSV format)", type=["csv"]) | |
if uploaded_file is not None: | |
filename = uploaded_file.name | |
namespace = filename.split('.')[0] | |
st.info("Dataset Processing Started...") | |
st.session_state.df, st.session_state.namespace = load_and_process_data(uploaded_file) | |
st.info("Dataset Processing Completed...") | |
if st.session_state.namespace: | |
query = st.text_input("Enter your query about the data (or type 'exit' to quit):") | |
if query.lower() != 'exit': | |
vec = st.session_state.model.encode(query) | |
result = None | |
result = st.session_state.index.query( | |
namespace=st.session_state.namespace, | |
vector=vec.tolist(), | |
top_k=5, | |
include_values=False | |
) | |
st.subheader("Query Results:") | |
if result is not None: | |
id = result['matches'][0]['id'] | |
data = st.session_state.df | |
answer = data[data['id'] == id]['Answer'].values[0] | |
st.write(answer) | |
if st.button("Delete Stored Data"): | |
st.session_state.index.delete(deleteAll=True, namespace =st.session_state.namespace) | |
st.stop() | |
if __name__ == "__main__": | |
main() | |