Added bigscience/bloom-petals

app.py

@@ -1,21 +1,26 @@
 import sys
-from transformers import AutoModelForCausalLM, AutoTokenizer
+
+import gradio as gr
 import torch
 import transformers
-import gradio as gr
+from transformers import AutoModelForCausalLM, AutoTokenizer
 
 sys.path.insert(0, './petals/')
 
-from src.client.remote_model import DistributedBloomForCausalLM
-
-MODEL_NAME = "bigscience/test-bloomd-6b3"  # select model you like
-INITIAL_PEERS = ["/ip4/193.106.95.184/tcp/31000/p2p/QmSg7izCDtowVTACbUmWvEiQZNY4wgCQ9T9Doo66K59X6q"]
+from petals.client.remote_model import DistributedBloomForCausalLM
 
-tokenizer_bloomd_6b3 = transformers.BloomTokenizerFast.from_pretrained("bigscience/test-bloomd-6b3")
-model_bloomd_6b3 = DistributedBloomForCausalLM.from_pretrained("bigscience/test-bloomd-6b3",
-                                                               initial_peers=INITIAL_PEERS,
+MODEL_NAME = "bigscience/test-bloomd-6b3"
+# INITIAL_PEERS = ["/ip4/193.106.95.184/tcp/31000/p2p/QmSg7izCDtowVTACbUmWvEiQZNY4wgCQ9T9Doo66K59X6q"]
+tokenizer_bloomd_6b3 = transformers.BloomTokenizerFast.from_pretrained(MODEL_NAME)
+model_bloomd_6b3 = DistributedBloomForCausalLM.from_pretrained(MODEL_NAME,
+                                                               # initial_peers=INITIAL_PEERS,
                                                                low_cpu_mem_usage=True, torch_dtype=torch.float32)
 
+MODEL_NAME = "bigscience/bloom-petals"
+tokenizer_bloomd = transformers.BloomTokenizerFast.from_pretrained(MODEL_NAME)
+model_bloomd = DistributedBloomForCausalLM.from_pretrained(MODEL_NAME,
+                                                           low_cpu_mem_usage=True, torch_dtype=torch.float32)
+
 
 tokenizer_DialoGPT_small = AutoTokenizer.from_pretrained("microsoft/DialoGPT-small")
 model_DialoGPT_small = AutoModelForCausalLM.from_pretrained("microsoft/DialoGPT-small")
@@ -47,8 +52,11 @@ def predict(
         model = model_DialoGPT_large
         tokenizer = tokenizer_DialoGPT_large
     elif model_name == 'test-bloomd-6b3':
-        model = tokenizer_bloomd_6b3
-        tokenizer = model_bloomd_6b3
+        model = model_bloomd_6b3
+        tokenizer = tokenizer_bloomd_6b3
+    elif model_name == 'bloom-petals':
+        model = model_bloomd
+        tokenizer = tokenizer_bloomd
     else:
         model = model_DialoGPT_medium
         tokenizer = tokenizer_DialoGPT_medium
@@ -81,7 +89,8 @@ gr.Interface(
                 'DialoGPT-small',
                 'DialoGPT-medium',
                 'DialoGPT-large',
-                'test-bloomd-6b3'
+                'test-bloomd-6b3',
+                'bloom-petals',
             ]
         ),
         gr.Radio(