From 4499469efbd7a3ceb2e39b589ce986c2362d9fe9 Mon Sep 17 00:00:00 2001
From: Aaron Pham <29749331+aarnphm@users.noreply.github.com>
Date: Sat, 18 Nov 2023 02:02:16 -0500
Subject: [PATCH] fix(annotations): check library through find_spec (#691)

Signed-off-by: Aaron <29749331+aarnphm@users.noreply.github.com>
---
 openllm-python/src/openllm/_llm.py  | 9 +++------
 openllm-python/src/openllm/_llm.pyi | 2 +-
 2 files changed, 4 insertions(+), 7 deletions(-)

diff --git a/openllm-python/src/openllm/_llm.py b/openllm-python/src/openllm/_llm.py
index 9493cc13..3b4807eb 100644
--- a/openllm-python/src/openllm/_llm.py
+++ b/openllm-python/src/openllm/_llm.py
@@ -1,5 +1,6 @@
 from __future__ import annotations
 import functools
+import importlib.util
 import logging
 import os
 import types
@@ -349,12 +350,8 @@ class LLM(t.Generic[M, T], ReprMixin):
 
   @property
   def adapter_map(self):
-    try:
-      import peft as _  # noqa: F401
-    except ImportError as err:
-      raise MissingDependencyError(
-        "Failed to import 'peft'. Make sure to do 'pip install \"openllm[fine-tune]\"'"
-      ) from err
+    if importlib.util.find_spec('peft') is None:
+      raise MissingDependencyError("Failed to import 'peft'. Make sure to do 'pip install \"openllm[fine-tune]\"'")
     if not self.has_adapters:
       raise AttributeError('Adapter map is not available.')
     assert self._adapter_map is not None
diff --git a/openllm-python/src/openllm/_llm.pyi b/openllm-python/src/openllm/_llm.pyi
index bac2b346..18e974d3 100644
--- a/openllm-python/src/openllm/_llm.pyi
+++ b/openllm-python/src/openllm/_llm.pyi
@@ -24,7 +24,7 @@ from openllm_core.utils.representation import ReprArgs
 from ._quantisation import QuantizationConfig
 from ._runners import Runner
 
-InjectedModel = Union[PeftModel | PeftModelForCausalLM | PeftModelForSeq2SeqLM]
+InjectedModel = Union[PeftModel, PeftModelForCausalLM, PeftModelForSeq2SeqLM]
 
 class IdentifyingParams(TypedDict):
   configuration: str