Streamlit + Readme.md

README.md

@@ -10,4 +10,24 @@ pinned: false
 short_description: Semantic Search engine with Faiss
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-
+
+### For local deployment run 
+```
+main.py
+```
+which will create swagger app with endpoints on [localhost:8084](http://127.0.0.1:8084/docs). First endpoint return the top k semanticaly most similar prompts with query prompt. Second endpoint returns all similarites with query (only applicable for very small datasets).
+
+```
+data_reader.py
+```
+creates data of various prompts for encoding into vector database. Local database encoded only 6000 prompts.
+
+Faiss index that is used is small and not optimized, used for experimental datasets. Search is brute force, not optimised. 
+
+### Streamlit
+```
+streamlit run app.py
+```
+should be run for streamlit app, it can be assessed locally on http://localhost:8501.
+

app.py

@@ -6,13 +6,22 @@ from models.Query import Query, SimilarPrompt, SearchResponse, PromptVector, Vec
 from sentence_transformers import SentenceTransformer
 import os
 
-# Path to your prompts data (you need to upload this file to your Hugging Face space)
-prompt_path = "models/prompts_data.jsonl"  # Update this to the correct path in your space
-
-# Initialize search engine and model
-prompts = load_prompts_from_jsonl(prompt_path)
-search_engine = PromptSearchEngine()
-search_engine.add_prompts_to_vector_database(prompts)
+# Cache the prompts data to avoid reloading every time
+@st.cache_data
+def load_prompts():
+    prompt_path = "models/prompts_data.jsonl"
+    return load_prompts_from_jsonl(prompt_path)
+
+# Cache the search engine initialization
+@st.cache_resource
+def get_search_engine():
+    search_engine = PromptSearchEngine()
+    prompts = load_prompts()
+    search_engine.add_prompts_to_vector_database(prompts)
+    return search_engine
+
+# Initialize search engine only once
+search_engine = get_search_engine()
 
 # Streamlit App Interface
 st.title("Prompt Search Engine")
@@ -27,17 +36,12 @@ k = st.number_input("Number of similar prompts to retrieve:", min_value=1, max_v
 # Button to trigger search
 if st.button("Search Prompts"):
     if query_input:
-        query = Query(prompt=query_input)
-        similar_prompts, distances = search_engine.most_similar(query.prompt, top_k=k)
+        similar_prompts, distances = search_engine.most_similar(query_input, top_k=k)
 
         # Format and display search results
-        response = [
-            SimilarPrompt(prompt=prompt, distance=float(distance))
-            for prompt, distance in zip(similar_prompts, distances)
-        ]
-        st.write("Search Results:")
-        for result in response:
-            st.write(f"Prompt: {result.prompt}, Distance: {result.distance}")
+        st.write(f"Search Results: ")
+        for i, (prompt, distance) in enumerate(zip(similar_prompts, distances)):
+            st.write(f"{i+1}. Prompt: {prompt}, Distance: {distance}")
     else:
         st.error("Please enter a prompt.")
 
@@ -47,54 +51,8 @@ st.write("### Vector Similarities")
 
 if st.button("Retrieve All Vector Similarities"):
     if query_input:
-        query = Query(prompt=query_input)
-        query_embedding = search_engine.model.encode([query.prompt])  # Encode the prompt to a vector
+        query_embedding = search_engine.model.encode([query_input])  # Encode the prompt to a vector
         all_similarities = search_engine.cosine_similarity(query_embedding, search_engine.index)
-
-        # Format and display vector similarities
-        response = [
-            PromptVector(vector=index, distance=float(distance))
-            for index, distance in enumerate(all_similarities)
-        ]
-        st.write("Vector Similarities:")
-        for result in response:
-            st.write(f"Vector Index: {result.vector}, Distance: {result.distance}")
+        st.write(f"Vector Similarities: {all_similarities}")
     else:
-        st.error("Please enter a prompt.")
-
-
-
-# # streamlit_app.py
-# import streamlit as st
-# import requests
-
-# # Streamlit app title
-# st.title("Top K Search with Vector DataBase")
-
-# # FastAPI endpoint URL
-# # url = "http://localhost:8084/search/"
-# url = "https://huggingface.co/search/"
-
-# # Input fields in Streamlit
-# id = st.text_input("Enter ID:", value="1")
-# prompt = st.text_input("Enter your prompt:")
-# k = st.number_input("Top K results:", min_value=1, max_value=100, value=3)
-
-# # Trigger the search when the button is clicked
-# if st.button("Search"):
-#     # Construct the request payload
-#     payload = {
-#         "id": id,
-#         "prompt": prompt,
-#         "k": k
-#     }
-
-#     # Make the POST request
-#     response = requests.post(url, json=payload)
-
-#     # Handle the response
-#     if response.status_code == 200:
-#         results = response.json()
-#         st.write(results)
-#     else:
-#         st.error(f"Error: {response.status_code} - {response.text}")
+        st.error("Please enter a prompt.")

main.py

@@ -30,7 +30,7 @@ def read_root():
 
 @app.post("/search/")
 async def search_prompts(query: Query, k: int = 3):
-    print(f'Prompt: {query.prompt}')
+    print(f'Prompt: {query}')
     similar_prompts, distances = search_engine.most_similar(query.prompt, top_k=k)
     print(f'Similar Prompts {similar_prompts}')
     print(f'Distances {distances}')
@@ -48,6 +48,9 @@ async def all_vectors(query: Query):
 
     query_embedding = search_engine.model.encode([query.prompt])  # Encode the prompt to a vector
     all_similarities = search_engine.cosine_similarity(query_embedding, search_engine.index)
+    print(f'Prompt: {query}')
+    print(f'All Vector Similarities: {all_similarities}')
+    print(40*'****')
     response = [
         PromptVector(vector=index, distance=float(distance)) 
         for index, distance in enumerate(all_similarities)
@@ -56,7 +59,7 @@ async def all_vectors(query: Query):
 
 if __name__ == "__main__":
     # Server Config
-    SERVER_HOST_IP = socket.gethostbyname(socket.gethostname())
-    # SERVER_HOST_IP = socket.gethostbyname("localhost") # for local deployment
+    # SERVER_HOST_IP = socket.gethostbyname(socket.gethostname())
+    SERVER_HOST_IP = socket.gethostbyname("localhost") # for local deployment
     SERVER_PORT = int(8084)
     uvicorn.run(app, host=SERVER_HOST_IP, port=SERVER_PORT)

models/data_reader.py

@@ -32,11 +32,13 @@ def read_data(jsonl_file_path):
             print(row)
 
 def load_prompts_from_jsonl(file_path):
+    print('Loading prompts from:', file_path)
     prompts = []
     with open(file_path, 'r') as f:
         for line in f:
             data = json.loads(line)  # Each line is a JSON object
             prompts.append(data)  # Extract the 'prompt' field
+    print("Data loaded successfully.")
     return prompts
 
 

models/prompt_search_engine.py

@@ -6,6 +6,7 @@ import faiss
 
 class PromptSearchEngine:
     def __init__(self, model_name='bert-base-nli-mean-tokens'):
+        print("Search engine started!")
         self.model = SentenceTransformer(model_name)
         # Initialize FAISS index with right number of dimensions
         self.embedding_dimension = self.model.get_sentence_embedding_dimension()
@@ -14,13 +15,16 @@ class PromptSearchEngine:
 
 
     def add_prompts_to_vector_database(self, prompts):
+        print("Data encoding started...")
         embeddings = self.model.encode(prompts)
         self.index.add(np.array(embeddings).astype('float32'))  
         self.prompts_track.extend(prompts)
+        print("Data encoding completed!")
 
 
     def most_similar(self, query, top_k=5):
-        # Encode the query
+        # Encode the 
+        print('Finding the most similar vectors')
         query_embedding = self.model.encode([query]).astype('float32')
         
         # Optimizovana pretraga ali moramo promeniti vrstu indeksa
@@ -37,7 +41,7 @@ class PromptSearchEngine:
             Args: query_vector: The query vector to compare against the corpus vectors. corpus_vectors: The set of corpus vectors to compare against the query vector. 
             Returns: The cosine similarity between the query vector and the corpus vectors.
             """
-     
+        print('Searching for all similarities...')
         query_vector = np.array(query_vector).astype('float32')
         query_norm = query_vector / np.linalg.norm(query_vector)