Use smaller model for tests with petals

Author: mudler

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,9 +65,9 @@ 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")
+                req = backend_pb2.PredictOptions(Prompt="The capital of France is")
                 resp = stub.Predict(req)
                 self.assertIsNotNone(resp.message)
         except Exception as err:

backend/python/vllm/test_backend_vllm.py

@@ -66,7 +66,7 @@ def test_text(self):
                 stub = backend_pb2_grpc.BackendStub(channel)
                 response = stub.LoadModel(backend_pb2.ModelOptions(Model="facebook/opt-125m"))
                 self.assertTrue(response.success)
-                req = backend_pb2.PredictOptions(prompt="The capital of France is")
+                req = backend_pb2.PredictOptions(Prompt="The capital of France is")
                 resp = stub.Predict(req)
                 self.assertIsNotNone(resp.message)
         except Exception as err: