Pravincoder commited on
Commit
48a0fce
Β·
verified Β·
1 Parent(s): 2eb0f74

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +40 -38
app.py CHANGED
@@ -1,53 +1,55 @@
1
  import streamlit as st
2
- import tempfile
3
-
4
- # Use DuckDB instead of SQLite for ChromaDB
5
- import chromadb
6
- from chromadb.config import Settings
7
- # Import your RAG components
8
- from rag import build_chroma_store, pre_processing_csv, ask_query
9
  from sentence_transformers import SentenceTransformer
 
 
10
 
11
- # Temporary directory for persistence in Streamlit Cloud
12
- temp_dir = tempfile.TemporaryDirectory()
 
 
 
13
 
 
14
  @st.cache_resource
15
- def load_data(csv_path):
16
- """Load and process data, caching the results."""
17
- docs, metas = pre_processing_csv(csv_path)
18
-
19
- # Use DuckDB + Parquet instead of SQLite
20
- client = chromadb.Client(Settings(
21
- persist_directory=temp_dir.name
22
- ))
23
-
24
- collection, model = build_chroma_store(docs, metas, client=client)
25
- return collection, model
26
-
27
- # Load your CSV data
28
- csv_path = "shl_products.csv" # Update path if needed
29
- collection, model = load_data(csv_path)
 
30
 
31
  # Streamlit UI
32
  st.title("🧠 RAG Model Query Interface")
33
- st.write("Enter a query to get relevant SHL test assessments.")
34
 
35
- # Query input
36
- user_query = st.text_input("Enter your query:")
37
 
38
- if st.button("Submit"):
39
  if user_query:
40
- results = ask_query(user_query, model, collection)
 
 
41
  if results:
42
- st.write(f"πŸ“Š Results for query: {user_query}")
43
- st.write("=" * 80)
44
  for i, (doc, meta) in enumerate(results, 1):
45
- st.markdown(f"πŸ”Ή **Result {i}**")
46
- st.markdown(f"πŸ§ͺ **Test Name:** {meta['Test Name']}")
47
- st.markdown(f"πŸ”— **Link:** [https://www.shl.com{meta['Test Link']}]")
48
- st.markdown(f"πŸ“„ **Chunk:** {doc}")
49
- st.write("-" * 80)
50
  else:
51
- st.warning("No results found.")
52
  else:
53
- st.warning("Please enter a query.")
 
1
  import streamlit as st
2
+ from rag import pre_processing_csv, build_pinecone_store, ask_query
 
 
 
 
 
 
3
  from sentence_transformers import SentenceTransformer
4
+ from dotenv import load_dotenv
5
+ import os
6
 
7
+ # Load environment variables from .env file
8
+ load_dotenv()
9
+ PINECONE_API_KEY = os.getenv("PINECONE_API_KEY")
10
+ PINECONE_ENV = os.getenv("PINECONE_ENV") # e.g., "us-west-2"
11
+ PINECONE_INDEX_NAME = os.getenv("PINECONE_INDEX_NAME", "shl-test-index")
12
 
13
+ # Cache data processing and vector store creation for efficiency
14
  @st.cache_resource
15
+ def load_data():
16
+ csv_path = "shl_products.csv" # Update the path if necessary
17
+ # Step 1: Preprocess the CSV and create document chunks
18
+ documents, metadatas = pre_processing_csv(csv_path)
19
+
20
+ # Load the SentenceTransformer model (this may take some time at first run)
21
+ model = SentenceTransformer("all-MiniLM-L6-v2")
22
+
23
+ # Step 2: Build the Pinecone vector store with embeddings
24
+ index, model, embeddings, documents, metadatas = build_pinecone_store(
25
+ documents, metadatas, model, PINECONE_INDEX_NAME, PINECONE_API_KEY, PINECONE_ENV
26
+ )
27
+ return index, model
28
+
29
+ # Load and cache the data, index, and model
30
+ index, model = load_data()
31
 
32
  # Streamlit UI
33
  st.title("🧠 RAG Model Query Interface")
34
+ st.write("Enter a query to retrieve relevant SHL assessments.")
35
 
36
+ # Query input widget
37
+ user_query = st.text_input("Your query:")
38
 
39
+ if st.button("Submit Query"):
40
  if user_query:
41
+ # Step 3: Query the RAG model using the Pinecone index
42
+ results = ask_query(user_query, model, index, k=10)
43
+
44
  if results:
45
+ st.markdown(f"### Results for: `{user_query}`")
 
46
  for i, (doc, meta) in enumerate(results, 1):
47
+ st.markdown(f"**Result {i}:**")
48
+ st.markdown(f"**Test Name:** {meta.get('Test Name', '')}")
49
+ st.markdown(f"**Test Link:** [SHL Link](https://www.shl.com{meta.get('Test Link', '')})")
50
+ st.markdown(f"**Chunk:** {doc}")
51
+ st.markdown("---")
52
  else:
53
+ st.info("No results found. Please try a different query.")
54
  else:
55
+ st.warning("Please enter a query to search.")