search_engine / app.py
Vitomir Jovanović
New Streamlit code for Hugging Face deployment
e9fda99
raw
history blame
3.44 kB
import streamlit as st
from models.vectorizer import Vectorizer
from models.prompt_search_engine import PromptSearchEngine
from models.data_reader import load_prompts_from_jsonl
from models.Query import Query, SimilarPrompt, SearchResponse, PromptVector, VectorResponse
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)
# Streamlit App Interface
st.title("Prompt Search Engine")
st.write("Search for similar prompts using the local search engine.")
# Input for the user's prompt
query_input = st.text_input("Enter your prompt:")
# Number of similar prompts to retrieve (k)
k = st.number_input("Number of similar prompts to retrieve:", min_value=1, max_value=10, value=3)
# 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)
# 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}")
else:
st.error("Please enter a prompt.")
# Additional functionality for vector similarity
st.write("---")
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
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}")
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}")