feat(server): support RefinedWeb models (#379)

huggingface · May 30, 2023 · b8b950b · b8b950b
1 parent bf7f1d5
commit b8b950b
Show file tree

Hide file tree

Showing 12 changed files with 3,236 additions and 27 deletions.
diff --git a/integration-tests/conftest.py b/integration-tests/conftest.py
@@ -205,7 +205,10 @@ def event_loop():
 def launcher(event_loop):
     @contextlib.contextmanager
     def local_launcher(
-        model_id: str, num_shard: Optional[int] = None, quantize: Optional[str] = None
+        model_id: str,
+        num_shard: Optional[int] = None,
+        quantize: Optional[str] = None,
+        trust_remote_code: bool = False,
     ):
         port = random.randint(8000, 10_000)
         master_port = random.randint(10_000, 20_000)
@@ -230,6 +233,9 @@ def local_launcher(
             args.extend(["--num-shard", str(num_shard)])
         if quantize:
             args.append("--quantize")
+            args.append("bitsandbytes")
+        if trust_remote_code:
+            args.append("--trust-remote-code")
 
         env = os.environ
         env["LOG_LEVEL"] = "info,text_generation_router=debug"
@@ -250,7 +256,10 @@ def local_launcher(
 
     @contextlib.contextmanager
     def docker_launcher(
-        model_id: str, num_shard: Optional[int] = None, quantize: Optional[str] = None
+        model_id: str,
+        num_shard: Optional[int] = None,
+        quantize: Optional[str] = None,
+        trust_remote_code: bool = False,
     ):
         port = random.randint(8000, 10_000)
 
@@ -260,6 +269,9 @@ def docker_launcher(
             args.extend(["--num-shard", str(num_shard)])
         if quantize:
             args.append("--quantize")
+            args.append("bitsandbytes")
+        if trust_remote_code:
+            args.append("--trust-remote-code")
 
         client = docker.from_env()