Spaces:
Running
Running
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}") | |