Added files

Dockerfile

@@ -0,0 +1,16 @@
+# read the doc: https://huggingface.co/docs/hub/spaces-sdks-docker
+# you will also find guides on how best to write your Dockerfile
+
+FROM python:3.9
+
+RUN useradd -m -u 1000 user
+USER user
+ENV PATH="/home/user/.local/bin:$PATH"
+
+WORKDIR /app
+
+COPY --chown=user ./requirements.txt requirements.txt
+RUN pip install --no-cache-dir --upgrade -r requirements.txt
+
+COPY --chown=user . /app
+CMD ["uvicorn", "run:app", "--host", "0.0.0.0", "--port", "7860"]

promptSearchEngine.py

@@ -0,0 +1,83 @@
+from typing import List, Sequence, Tuple
+import numpy as np
+from vectorizer import Vectorizer
+
+def cosine_similarity(
+        query_vector: np.ndarray,
+        corpus_vectors: np.ndarray
+    )-> np.ndarray:
+
+        """Calculate cosine similarity between prompt vectors.
+        Args:
+        query_vector: Vectorized prompt query of shape (1, D).
+        corpus_vectors: Vectorized prompt corpus of shape (N, D).
+        Returns: The vector of shape (N,) with values in range [-1, 1] where 1
+        is max similarity i.e., two vectors are the same.
+        """
+        dot_product = np.dot( corpus_vectors, query_vector)
+        magnitude_A = np.linalg.norm(corpus_vectors, axis=1)
+        magnitude_B = np.linalg.norm(query_vector)
+
+        cosine_sim = dot_product / (magnitude_A * magnitude_B)
+        return np.around(cosine_sim, 4)
+        # return np.format_float_positional(cosine_sim, precision = 4)
+
+
+class PromptSearchEngine:
+    def __init__(self, prompts: Sequence[str], model) -> None:
+        """Initialize search engine by vectorizing prompt corpus.
+        Vectorized prompt corpus should be used to find the top n most
+        similar prompts w.r.t. user’s input prompt.
+        Args:
+        prompts: The sequence of raw prompts from the dataset.
+        """
+        self.prompts = prompts
+        self.vectorizer = Vectorizer(model)
+        self.corpus_embeddings = self.vectorizer.transform(prompts)
+    def most_similar(
+    self,
+    query: str,
+    n: int = 5
+    ) -> List[Tuple[float, str]]:
+        """Return top n most similar prompts from corpus.
+        Input query prompt should be vectorized with chosen Vectorizer.
+        After
+        that, use the cosine_similarity function to get the top n most
+        similar
+        prompts from the corpus.
+        Args:
+        query: The raw query prompt input from the user.
+        n: The number of similar prompts returned from the corpus.
+        Returns:
+        The list of top n most similar prompts from the corpus along
+        with similarity scores. Note that returned prompts are
+        verbatim.
+        """
+        most_similar_prompts = []
+        prompt_embedding = self.vectorizer.transform([query]).flatten()
+        corpus_embeddings = self.corpus_embeddings
+
+        result = cosine_similarity(prompt_embedding, corpus_embeddings)
+        
+        for i in range(len(self.prompts)):
+            most_similar_prompts.append((result[i], self.prompts[i]))
+
+        prompt_score_sorted = sorted(most_similar_prompts, key=lambda x: x[0], reverse=True)
+
+        return prompt_score_sorted[0:n]
+    def display_prompts(self, prompts):
+        """Display the list of prompts with their similarity scores."""
+        if prompts:
+            for i, (score, prompt) in enumerate(prompts, 1):
+                print(f"{i}. {prompt} (Similarity: {score:.4f})")
+        else:
+            print("No prompts found.")
+    def stringify_prompts(self, prompts):
+        """Save the list of prompts with their similarity scores."""
+        strings = []
+        if prompts:
+            for i, (score, prompt) in enumerate(prompts, 1):
+                strings.append(f"{i}. {prompt} (Similarity: {score:.4f})")
+            return strings
+        else:
+            return []

run.py

@@ -0,0 +1,46 @@
+import json
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel
+from promptSearchEngine import PromptSearchEngine
+from vectorizer import Vectorizer
+from datasets import load_dataset
+from sentence_transformers import SentenceTransformer
+
+EMBEDDING_MODEL = "all-MiniLM-L6-v2"
+DATASET = "Gustavosta/Stable-Diffusion-Prompts"
+
+
+    
+model = SentenceTransformer(EMBEDDING_MODEL)
+dataset = load_dataset(DATASET , split="test[:1%]")
+promptSearchEngine = PromptSearchEngine(dataset["Prompt"], model)
+
+class SearchRequest(BaseModel):
+    query: str 
+    n: int | None = 5
+
+app = FastAPI()
+
+@app.get("/")
+async def root():
+    return {"message": 'GET /docs'}
+
+@app.get("/search")
+async def search(q: str, n: int = 5):
+    results = []
+    if q.isspace() or q =="":
+        return {"message": "Enter query"}
+    else:
+        results = promptSearchEngine.most_similar(q, n)
+    if not results:
+        raise HTTPException(status_code=404, detail="No prompts found.")
+    return promptSearchEngine.stringify_prompts(results)
+
+
+@app.post("/search")
+async def searchPost(request: SearchRequest):
+    results = promptSearchEngine.most_similar(request.query, request.n)
+    if not results:
+        raise HTTPException(status_code=404, detail="No prompts found.")
+    formatted_results = [{"similarity": float(similarity), "prompt": prompt } for similarity, prompt in results]
+    return { "data" : formatted_results }

run_local_ui.py

@@ -0,0 +1,67 @@
+from pydantic import BaseModel
+from promptSearchEngine import PromptSearchEngine
+from datasets import load_dataset
+from sentence_transformers import SentenceTransformer
+import streamlit as st
+
+EMBEDDING_MODEL = "all-MiniLM-L6-v2"
+DATASET = "Gustavosta/Stable-Diffusion-Prompts"
+
+class SearchRequest(BaseModel):
+    query: str 
+    n: int | None = 5
+    
+# model = SentenceTransformer("all-MiniLM-L6-v2")
+# dataset = load_dataset("Gustavosta/Stable-Diffusion-Prompts" , split="test[:1%]")
+# promptSearchEngine = PromptSearchEngine(dataset["Prompt"], model)
+
+@st.cache_resource  
+def load_model():
+    """Initialize pretrained model for vectorizing.
+        @st.cache_resource anotation enables caching for Streamlit.
+    """
+    return SentenceTransformer(EMBEDDING_MODEL)
+
+@st.cache_resource  
+def load_dataSet():
+    """Initialize pretrained model for vectorizing.
+        @st.cache_resource anotation enables caching for Streamlit. 
+    """
+    return load_dataset(DATASET , split="test[:1%]")
+
+@st.cache_resource  
+def load_searchEngine(prompts, _model):
+    """Initialize search engine and vectorize raw propmpts from dataset.
+        @st.cache_resource anotation enables caching for Streamlit.
+        Args:
+        prompts: The sequence of raw prompts from the dataset.
+        model: The model for vectorizing.
+    """
+    return PromptSearchEngine(prompts, _model)
+
+model = load_model()
+dataset = load_dataSet()
+promptSearchEngine = load_searchEngine(dataset["Prompt"], model)
+
+
+with st.form("search_form"):
+    st.write("Prompt Search Engine")
+    query = st.text_area("Prompt to search")
+    number = st.number_input("Number of similar prompts", value = 5, min_value=0, max_value=100)
+    submitted = st.form_submit_button("Submit")
+    if submitted:
+        result = promptSearchEngine.most_similar(query, number)
+        st.dataframe(
+            result, 
+            use_container_width=True,
+            column_config={
+                1: st.column_config.NumberColumn(
+                    "Similarity", 
+                    help="Range in [-1, 1] where 1 is max similarity, means that prompts are identical.",
+                    format= "%.4f"
+                ),
+                2: st.column_config.TextColumn("Prompts", help="The simlar prompts"),
+            },
+        )
+
+

run_ui.py

@@ -0,0 +1,33 @@
+from pydantic import BaseModel
+from promptSearchEngine import PromptSearchEngine
+from datasets import load_dataset
+from sentence_transformers import SentenceTransformer
+import streamlit as st
+import requests
+import json
+
+st.title('Prompt Search Engine')
+
+with st.form("search_form"):
+    st.write("Prompt Search Engine")
+    query = st.text_area("Prompt to search")
+    number = st.number_input("Number of similar prompts", value = 5, min_value=0, max_value=100)
+    submitted = st.form_submit_button("Submit")
+    if submitted:
+        inputs = {"query": query, "n": number}
+        result = requests.post(url = "http://localhost:8000/search", data = json.dumps(inputs))
+        result = result.json()
+        st.dataframe(
+            result["data"], 
+            use_container_width=True,
+            column_config={
+                "similarity": st.column_config.NumberColumn(
+                    "Similarity", 
+                    help="Range in [-1, 1] where 1 is max similarity, means that prompts are identical.",
+                    format= "%.4f"
+                ),
+                "prompt": st.column_config.TextColumn("Prompts", help="The simlar prompts"),
+            },
+        )
+
+

vectorizer.py

@@ -0,0 +1,24 @@
+from typing import Sequence
+import numpy as np
+
+
+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
+
+    def transform(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 as a numpy array."""
+        vectorized = self.model.encode(prompts, show_progress_bar=True)
+        return vectorized