From bee1224aa2e6bd5202afcf9f34d9bf99c139a895 Mon Sep 17 00:00:00 2001
From: ishaan-jaff <ishaanjaffer0324@gmail.com>
Date: Tue, 19 Sep 2023 09:16:05 -0700
Subject: [PATCH] dev4

---
 litellm/llms/petals.py | 10 +++-------
 pyproject.toml         |  2 +-
 2 files changed, 4 insertions(+), 8 deletions(-)

diff --git a/litellm/llms/petals.py b/litellm/llms/petals.py
index 648834540e..e55a643952 100644
--- a/litellm/llms/petals.py
+++ b/litellm/llms/petals.py
@@ -33,8 +33,8 @@ def completion(
     model = model
 
     tokenizer = AutoTokenizer.from_pretrained(model, use_fast=False, add_bos_token=False)
-    model = AutoDistributedModelForCausalLM.from_pretrained(model)
-    model = model.cuda()
+    model_obj = AutoDistributedModelForCausalLM.from_pretrained(model)
+    model_obj = model_obj.cuda()
 
     prompt = ""
     for message in messages:
@@ -60,7 +60,7 @@ def completion(
     
     ## COMPLETION CALL
     inputs = tokenizer(prompt, return_tensors="pt")["input_ids"].cuda()
-    outputs = model.generate(inputs, max_new_tokens=5)
+    outputs = model_obj.generate(inputs, max_new_tokens=5)
     print(outputs)
 
 
@@ -71,14 +71,10 @@ def completion(
             original_response=outputs,
             additional_args={"complete_input_dict": optional_params},
         )
-    print_verbose(f"raw model_response: {outputs}")
     ## RESPONSE OBJECT
     output_text = tokenizer.decode(outputs[0])
-    print("output text")
-    print(output_text)
     model_response["choices"][0]["message"]["content"] = output_text
 
-    ## CALCULATING USAGE - baseten charges on time, not tokens - have some mapping of cost here. 
     prompt_tokens = len(
         encoding.encode(prompt)
     ) 
diff --git a/pyproject.toml b/pyproject.toml
index 5a679fa184..dac1ed4712 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "litellm"
-version = "0.1.700"
+version = "0.1.700.dev4"
 description = "Library to easily interface with LLM API providers"
 authors = ["BerriAI"]
 license = "MIT License"