diff --git a/rerankers/models/ranker.py b/rerankers/models/ranker.py
index 6c8f45a..bf7f326 100644
--- a/rerankers/models/ranker.py
+++ b/rerankers/models/ranker.py
@@ -30,8 +30,8 @@ def rank(
     async def rank_async(
         self,
         query: str,
-        docs: List[str],
-        doc_ids: Optional[Union[List[str], str]] = None,
+        docs: Union[str, List[str], Document, List[Document]],
+        doc_ids: Optional[Union[List[str], List[int]]] = None,
     ) -> RankedResults: