diff --git a/inference/download_model.py b/inference/download_model.py
index e8443e1..5ffee0e 100644
--- a/inference/download_model.py
+++ b/inference/download_model.py
@@ -37,8 +37,7 @@
 )
 def download_model(model_name, model_revision, force_download=False):
     from huggingface_hub import snapshot_download
-    from transformers import AutoTokenizer, AutoModelForCausalLM
-    import json
+    from transformers import AutoTokenizer
     import os
 
     volume.reload()
diff --git a/inference/serving-non-optimized-fastapi.py b/inference/serving-non-optimized-fastapi.py
index 0c55876..1c69be8 100644
--- a/inference/serving-non-optimized-fastapi.py
+++ b/inference/serving-non-optimized-fastapi.py
@@ -53,9 +53,6 @@ def load(self):
 
     @modal.method()
     def generate(self, chat):
-        import torch
-        from torch.nn.functional import softmax
-
         tokenizer = self.tokenizer
         model = self.model
 
diff --git a/pyproject.toml b/pyproject.toml
index ea8eabe..a5c5ab8 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,5 +1,5 @@
 [project]
-name = "llamaguard-7b
+name = "llamaguard-7b"
 version = "0.0.1"
 description = "Template repo for Guardrails Hub validators."
 authors = [
@@ -29,5 +29,7 @@ testpaths = [
 [tool.pyright]
 include = ["validator"]
 
-[tool.setuptools.packages.find]
-include = ["validator", "inference", "validator.*", "inference.*"]
\ No newline at end of file
+[tool.setuptools]
+packages = [
+    "validator"
+]
\ No newline at end of file