update README

README.md

@@ -9,4 +9,84 @@ license: apache-2.0
 app_port: 8000
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+## Overview
+
+The **Prompt Engine** project is a take-home assignment for the SmartCat recruitment process. This project comprises two main components:
+
+- **Backend**: A REST API that provides functionality to find similar prompts.
+- **Frontend**: A user interface for interacting with the backend service.
+
+### Service Description
+
+The service retrieves and displays the 5 most similar prompts to a user’s query. This helps improve user queries by showing related examples from a database of well-curated prompts.
+
+## Setup Instructions
+
+### Prerequisites
+
+Before you start, ensure you have the following tools installed:
+
+- `Poetry`
+- `Docker`
+
+### Installation
+
+1. **Install Dependencies**:
+
+   ```bash
+   poetry install
+   ```
+
+2. **Activate the Virtual Environment**:
+
+   ```bash
+   poetry shell
+   ```
+
+### Backend Details
+
+The backend API provides the following endpoint:
+
+- **Endpoint**: POST /most_similar
+- **Description**: Accepts a user query and the number of similar prompts to return.
+
+### Docker
+
+The repository includes a Dockerfile to build and run the backend service. To build the image:
+
+```bash
+docker build -t prompt-search-engine .
+```
+
+To start the backend service, run:
+
+```bash
+docker run -d -p 8000:8000 prompt-serach-engine
+```
+
+This command will build the Docker image and start the service.
+
+### Deployment
+
+The project is integrated with HuggingFace Spaces. The solution is automatically deployed to HuggingFace with every push to the repository.
+
+In case you want to set up your own HuggingFace Space, you must create a HF token and add it with:
+
+```
+git remote set-url origin https://USERNAME:[email protected]/spaces/USERNAME/REPO_NAME.git
+```
+
+### User Interface
+
+The frontend UI is simple and includes:
+
+- A text field for entering the user query.
+- A slider to specify the number of similar prompts to retrieve
+
+To start only the UI service you can run:
+
+```bash
+poe frontend --api_url http://localhost:8000
+```
+
+The default api points to HF Space of this repository.

frontend/app_ui.py

@@ -1,13 +1,27 @@
-import streamlit as st
+import argparse
+
 import requests
+import streamlit as st
 
-# Define the URL of the FastAPI service
-API_URL = "https://lazarr19-prompt-engine.hf.space/most_similar"  # Adjust if the FastAPI service is hosted elsewhere
+
+def parse_arguments():
+    """Parse command-line arguments."""
+    parser = argparse.ArgumentParser(description="Prompt Similarity Finder")
+    parser.add_argument(
+        "--api_url",
+        type=str,
+        default="https://lazarr19-prompt-engine.hf.space",
+        help="The URL of the FastAPI service",
+    )
+    return parser.parse_args()
 
 
-def get_similar_prompts(query, n):
+def get_similar_prompts(api_url, query, n):
+    """Fetch similar prompts from the FastAPI service."""
     try:
-        response = requests.post(API_URL, json={"query": query, "n": n})
+        response = requests.post(
+            f"{api_url}/most_similar", json={"query": query, "n": n}
+        )
         response.raise_for_status()  # Raise an exception for HTTP errors
         return response.json()
     except requests.RequestException as e:
@@ -16,6 +30,7 @@ def get_similar_prompts(query, n):
 
 
 def get_color(score):
+    """Determine the color based on the score."""
     if score >= 0.8:
         return "green"
     elif score >= 0.5:
@@ -24,7 +39,8 @@ def get_color(score):
         return "red"
 
 
-def main():
+def main(api_url):
+    """Main function to run the Streamlit app."""
     st.title("Prompt Similarity Finder")
 
     # User input for query
@@ -36,7 +52,7 @@ def main():
     if st.button("Find Similar Prompts"):
         if query:
             with st.spinner("Fetching similar prompts..."):
-                result = get_similar_prompts(query, n)
+                result = get_similar_prompts(api_url, query, n)
                 if result:
                     similar_prompts = result.get("similar_prompts", [])
                     if similar_prompts:
@@ -57,4 +73,5 @@ def main():
 
 
 if __name__ == "__main__":
-    main()
+    args = parse_arguments()
+    main(args.api_url)

main.py

@@ -1,9 +1,11 @@
+from typing import List
+
 from fastapi import FastAPI, HTTPException
 from fastapi.responses import HTMLResponse
 from pydantic import BaseModel
-from typing import List
-from src.search_engine import PromptSearchEngine
+
 from src.prompt_loader import PromptLoader
+from src.search_engine import PromptSearchEngine
 
 # Constants
 SEED = 42

poe/common-tasks.toml

@@ -51,15 +51,6 @@ cmd  = """
 help     = "Run all checks on the code base"
 sequence = ["style", "types", "lint", "clean"]
 
-[tool.poe.tasks.ui]
+[tool.poe.tasks.frontend]
 help = "Start the UI"
-cmd  = "streamlit run ./frontend/app_ui.py"
-
-[tool.poe.tasks.build]
-help = "Build the Docker image"
-cmd  =  "docker build -t prompt-search-engine ."
-
-[tool.poe.tasks.start]
-help = "Run the Docker container"
-cmd  =  "docker run -d -p 8000:8000 prompt-serach-engine"
-
+cmd  = "streamlit run ./frontend/app_ui.py"

pyproject.toml

@@ -1,5 +1,5 @@
 [tool.poetry]
-name = "smart-cat-assignment-backend"
+name = "smart-cat-assignment"
 version = "0.0.1"
 description = "SmartCat Assignment"
 authors     = ["Lazar Radojevic <[email protected]>"]

src/prompt_loader.py

@@ -1,6 +1,7 @@
-from typing import Optional, List
-from datasets import load_dataset
 import random
+from typing import List, Optional
+
+from datasets import load_dataset
 
 
 class PromptLoader:

src/search_engine.py

@@ -2,8 +2,8 @@ from typing import List, Sequence, Tuple
 
 from sentence_transformers import SentenceTransformer
 
-from src.vectorizer import Vectorizer
 from src.similarity_scorer import SimilarityScorer
+from src.vectorizer import Vectorizer
 
 
 class PromptSearchEngine: