Use smaller model for tests with petals

mudler · Dec 20, 2023 · 5f56152 · 5f56152
1 parent 739ff29
commit 5f56152
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/backend/python/petals/test_petals.py b/backend/python/petals/test_petals.py
@@ -47,7 +47,8 @@ def test_load_model(self):
             self.setUp()
             with grpc.insecure_channel("localhost:50051") as channel:
                 stub = backend_pb2_grpc.BackendStub(channel)
-                response = stub.LoadModel(backend_pb2.ModelOptions(Model="petals-team/StableBeluga"))
+                response = stub.LoadModel(backend_pb2.ModelOptions(Model="bigscience/bloom-560m"))
+                print(response)
                 self.assertTrue(response.success)
                 self.assertEqual(response.message, "Model loaded successfully")
         except Exception as err:
@@ -64,7 +65,7 @@ def test_text(self):
             self.setUp()
             with grpc.insecure_channel("localhost:50051") as channel:
                 stub = backend_pb2_grpc.BackendStub(channel)
-                response = stub.LoadModel(backend_pb2.ModelOptions(Model="petals-team/StableBeluga"))
+                response = stub.LoadModel(backend_pb2.ModelOptions(Model="bigscience/bloom-560m"))
                 self.assertTrue(response.success)
                 req = backend_pb2.PredictOptions(prompt="The capital of France is")
                 resp = stub.Predict(req)