fix: :memo: Refactor docs to fix problems

.python-version

@@ -1 +0,0 @@
-3.11

Dockerfile

@@ -30,8 +30,8 @@ COPY requirements.txt ./
 COPY src/ ./src/
 COPY embeddings/ ./embeddings/
 
-RUN pip3 install --user -r requirements.txt
 ENV PATH=$HOME/.local/bin:$PATH
+RUN pip3 install --user -r requirements.txt
 
 EXPOSE 8501
 

requirements.txt

@@ -1,12 +1,5 @@
-numpy==1.26.4
-packaging==23.2
-setuptools==69.2.0
-fastapi==0.109.2
-uvicorn==0.27.1
-python-multipart==0.0.9
-requests==2.31.0
-torch==2.2.1
-torchvision==0.17.1
-Pillow==10.2.0
-streamlit==1.32.0 
+pillow
+requests
+streamlit
+torch
 transformers

src/similarity.py

@@ -35,7 +35,6 @@ device = "cuda" if torch.cuda.is_available() else "cpu"
 # --- Load Pretrained Model ---
 def get_model() -> ViTModel:
     """
-    TODO: Implement model loading
     - Load a pretrained model (e.g., ResNet18)
     - Remove the classification head
     - Set the model to evaluation mode
@@ -44,6 +43,10 @@ def get_model() -> ViTModel:
     Returns:
         torch.nn.Module: The prepared model
     """
+    # Ensure cache directory exists
+    cache_dir = os.getenv('TORCH_HOME', os.path.expanduser('~/.cache/torch'))
+    os.makedirs(cache_dir, exist_ok=True)
+    transformers_logger.info(f"Using cache directory: {cache_dir}")
 
     model = ViTModel.from_pretrained('imjeffhi/pokemon_classifier').to(device)
 
@@ -55,7 +58,6 @@ def get_model() -> ViTModel:
 
 
 # --- Image Preprocessing ---
-# TODO: Define your image transformation pipeline
 # Hint: Consider resizing, normalization, and tensor conversion
 transform = ViTImageProcessor.from_pretrained(get_model().name_or_path)
 
@@ -66,7 +68,6 @@ transform = ViTImageProcessor.from_pretrained(get_model().name_or_path)
 class PokemonSimilarity:
     def __init__(self, suppress_init_logs: bool = True) -> None:
         """
-        TODO: Initialize the similarity engine
         - Load the model
         - Load the database of Pokemon embeddings
         """
@@ -92,7 +93,6 @@ class PokemonSimilarity:
 
     def _load_db(self) -> dict | None:
         """
-        TODO: Implement database loading
         - Look for the embeddings file in different possible locations
         - Load the pickle file containing Pokemon embeddings
         - Handle cases where the file is not found