Search Engine

api.py

@@ -0,0 +1,33 @@
+# 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://127.0.0.1:8000/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}")

main.py

@@ -0,0 +1,58 @@
+import sys
+import os
+import copy
+import uvicorn
+import socket
+import logging
+import datetime
+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, Query_Multiple, SearchResponse, SimilarPrompt
+from decouple import config
+from fastapi import FastAPI, HTTPException, Depends, Body
+from sentence_transformers import SentenceTransformer
+
+
+
+prompt_path = r"C:\Users\jov2bg\Desktop\PromptSearch\models\prompts_data.jsonl"
+
+
+app = FastAPI(title="Search Prompt Engine", description="API for prompt search", version="1.0")
+
+prompts = load_prompts_from_jsonl(prompt_path)
+search_engine = PromptSearchEngine()
+search_engine.add_prompts_to_vector_database(prompts)
+
+@app.get("/")
+def read_root():
+    return {"message": "Prompt Search Engine is running!"}
+
+@app.post("/search/")
+async def search_prompts(query: Query, k: int = 3):
+    print(f'Prompt: {query.prompt}')
+    similar_prompts, distances = search_engine.most_similar(query.prompt, top_k=k)
+    print(f'Similar Prompts {similar_prompts}')
+    print(f'Distances {distances}')
+    print(40*'****')
+    # Format the response
+    response = [
+        SimilarPrompt(prompt=prompt, distance=float(distance)) 
+        for prompt, distance in zip(similar_prompts, distances)
+    ]
+    
+    return SearchResponse(results=response)
+
+@app.post("/all_vectors_similarities/")
+async def all_vectors(query: Query):
+
+    all_similarities = search_engine.cosine_similarity(query.prompt, search_engine.index)
+    response = [
+        SimilarPrompt(prompt=prompt, distance=float(distance)) 
+        for prompt, distance in all_similarities.items()
+    ]
+    return SearchResponse(results=response)
+
+if __name__ == "__main__":
+    localhost = socket.gethostbyname("localhost") 
+    uvicorn.run(app, host=localhost, port=8000)

models/Query.py

@@ -0,0 +1,20 @@
+from pydantic import BaseModel
+from typing import List
+
+
+
+class Query(BaseModel):
+    
+    id: str
+    prompt: str
+
+class Query_Multiple(BaseModel):
+    prompt: List[Query]
+
+
+class SimilarPrompt(BaseModel):
+    prompt: str
+    distance: float  
+
+class SearchResponse(BaseModel):
+    results: List[SimilarPrompt]

models/data_reader.py

@@ -0,0 +1,48 @@
+from datasets import load_dataset
+import json
+
+
+# Load the dataset
+base_url = "https://huggingface.co/datasets/jackyhate/text-to-image-2M/resolve/main/data_512_2M/data_{i:06d}.tar"
+num_shards = 46  # Number of webdataset tar files
+
+
+
+def download_data(base_url, num_shards):
+    # Download the data
+    urls = [base_url.format(i=i) for i in range(num_shards)]
+    dataset = load_dataset("webdataset", data_files={"train": urls}, split="train", streaming=True)
+    return dataset
+    
+def extract_prompts(dataset, json_file_path):
+    # Write data to the jsonl file
+    prompts = {}
+    with open(jsonl_file_path, 'w') as f:
+        for index, row in enumerate(dataset):
+            prompts[index] = row['json']['prompt']
+            f.write(json.dumps(prompts[index]) + '\n')
+
+
+def read_data(jsonl_file_path):
+
+    # Read data from the jsonl file
+    with open(jsonl_file_path, 'r') as f:
+        for line in f:
+            row = json.loads(line)
+            print(row)
+
+def load_prompts_from_jsonl(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
+    return prompts
+
+
+if __name__ == "__main__":
+    jsonl_file_path = r"C:\Users\jov2bg\Desktop\PromptSearch\models\prompts_data.jsonl"
+    num_shards = 1
+    dataset = download_data(num_shards, base_url)
+    extract_prompts(dataset, jsonl_file_path)
+    read_data(jsonl_file_path)

models/prompt_search_engine.py

@@ -0,0 +1,48 @@
+from typing import Sequence, List, Tuple
+from models.vectorizer import Vectorizer
+import numpy as np
+from sentence_transformers import SentenceTransformer
+import faiss
+
+class PromptSearchEngine:
+    def __init__(self, model_name='bert-base-nli-mean-tokens'):
+        self.model = SentenceTransformer(model_name)
+        # Initialize FAISS index with right number of dimensions
+        self.embedding_dimension = self.model.get_sentence_embedding_dimension()
+        self.index = faiss.IndexFlatL2(self.embedding_dimension)  # Euclidian distance index - brute force for small datasets
+        self.prompts_track = []  # To keep track of original prompts for returning results
+
+
+    def add_prompts_to_vector_database(self, prompts):
+        embeddings = self.model.encode(prompts)
+        self.index.add(np.array(embeddings).astype('float32'))  
+        self.prompts_track.extend(prompts)
+
+
+    def most_similar(self, query, top_k=5):
+        # Encode the query
+        query_embedding = self.model.encode([query]).astype('float32')
+        
+        # Optimizovana pretraga ali moramo promeniti vrstu indeksa
+        distances, indices = self.index.search(query_embedding, top_k)
+        
+        # Retrieve the corresponding prompts for the found indices
+        similar_prompts = [self.prompts_track[idx] for idx in indices[0]]
+        
+        return similar_prompts, distances[0]  # Return both the similar prompts and their distances
+
+
+    def cosine_similarity(query_vector: np.ndarray, corpus_vectors: np.ndarray) -> np.ndarray:
+        """Compute the cosine similarity between a query vector and a set of corpus vectors.
+        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.
+        """
+        similarities = {}
+        for index, vector in enumerate(corpus_vectors):
+            if np.linalg.norm(vector) == 0:
+                raise ValueError("One of the corpus vectors has zero norm.")
+            cos_similarity = np.dot(vector, query_vector) / (np.linalg.norm(vector) * np.linalg.norm(query_vector))
+            similarities[index] = cos_similarity
+        return similarities
+    
+      

models/vectorizer.py

@@ -0,0 +1,33 @@
+
+from sentence_transformers import SentenceTransformer
+import numpy as np
+from typing import Sequence
+import faiss
+
+
+
+
+
+class Vectorizer:
+    def __init__(self, model) -> None:
+        """Initialize the vectorizer with a pre-trained embedding model.
+        Args: model: The pre-trained embedding model to use for transforming prompts.
+        """
+        self.model = model
+        self.index_size = 50000
+        self.index = faiss.IndexFlatIP(self.index_size)
+        self.cached_index_idx_to_retrieval_db_idx = []
+    
+
+    def transform_and_add_to_index(self, prompts: Sequence[str]) -> np.ndarray:  
+        """Transform texts into numerical vectors using the specified model.
+        Args: prompts: The sequence of raw corpus prompts. Returns: Vectorized prompts
+        """
+        embeddings = self.model.encode(prompts)
+        embedding_dimension = embeddings.shape[1]
+        print('Embedding dimension:', embedding_dimension)
+        
+        self.index.add(np.array(embeddings))
+
+
+