From 946c81626df4694726626cd01e2c24ccbbc7b937 Mon Sep 17 00:00:00 2001
From: ishaan-jaff <ishaanjaffer0324@gmail.com>
Date: Tue, 19 Sep 2023 09:27:40 -0700
Subject: [PATCH] fix petals imports

---
 litellm/llms/petals.py | 11 ++++++++---
 1 file changed, 8 insertions(+), 3 deletions(-)

diff --git a/litellm/llms/petals.py b/litellm/llms/petals.py
index e55a6439526..13244cce646 100644
--- a/litellm/llms/petals.py
+++ b/litellm/llms/petals.py
@@ -26,9 +26,14 @@ def completion(
     litellm_params=None,
     logger_fn=None,
 ):
-    import torch
-    from transformers import AutoTokenizer
-    from petals import AutoDistributedModelForCausalLM
+    try:
+        import torch
+        from transformers import AutoTokenizer
+        from petals import AutoDistributedModelForCausalLM
+    except:
+        raise Exception(
+            "Importing torch, transformers, petals failed\nTry pip installing petals \npip install git+https://github.com/bigscience-workshop/petals"
+        )
 
     model = model