Repository: xorbitsai/inference

Issue:
ENH: including model revision into model spec
### Is your feature request related to a problem? Please describe
HF models & scripts can be updated, which may result in model re-downloading or introduce incompatible changes with xinference.

### Describe the solution you'd like
Add a model revison to each pytorch model spec. When caching a model, specify the model revision.




Discussion:

Resolving patch:
diff --git a/xinference/core/worker.py b/xinference/core/worker.py
--- a/xinference/core/worker.py
+++ b/xinference/core/worker.py
@@ -108,6 +108,7 @@ def _to_llm_description(
             "model_format": llm_spec.model_format,
             "model_size_in_billions": llm_spec.model_size_in_billions,
             "quantization": quantization,
+            "revision": llm_spec.model_revision,
         }
 
     @log_sync(logger=logger)
diff --git a/xinference/model/llm/llm_family.py b/xinference/model/llm/llm_family.py
--- a/xinference/model/llm/llm_family.py
+++ b/xinference/model/llm/llm_family.py
@@ -34,6 +34,7 @@ class GgmlLLMSpecV1(BaseModel):
     model_id: str
     model_file_name_template: str
     model_uri: Optional[str]
+    model_revision: Optional[str]
 
 
 class PytorchLLMSpecV1(BaseModel):
@@ -42,6 +43,7 @@ class PytorchLLMSpecV1(BaseModel):
     quantizations: List[str]
     model_id: str
     model_uri: Optional[str]
+    model_revision: Optional[str]
 
 
 class PromptStyleV1(BaseModel):
@@ -139,6 +141,7 @@ def cache_from_huggingface(
         assert isinstance(llm_spec, PytorchLLMSpecV1)
         huggingface_hub.snapshot_download(
             llm_spec.model_id,
+            revision=llm_spec.model_revision,
             local_dir=cache_dir,
             local_dir_use_symlinks=True,
         )
@@ -147,6 +150,7 @@ def cache_from_huggingface(
         file_name = llm_spec.model_file_name_template.format(quantization=quantization)
         huggingface_hub.hf_hub_download(
             llm_spec.model_id,
+            revision=llm_spec.model_revision,
             filename=file_name,
             local_dir=cache_dir,
             local_dir_use_symlinks=True,
diff --git a/xinference/model/llm/pytorch/core.py b/xinference/model/llm/pytorch/core.py
--- a/xinference/model/llm/pytorch/core.py
+++ b/xinference/model/llm/pytorch/core.py
@@ -47,7 +47,7 @@ class PytorchGenerateConfig(TypedDict, total=False):
 
 
 class PytorchModelConfig(TypedDict, total=False):
-    revision: str
+    revision: Optional[str]
     device: str
     gpus: Optional[str]
     num_gpus: int
@@ -80,7 +80,7 @@ def _sanitize_model_config(
     ) -> PytorchModelConfig:
         if pytorch_model_config is None:
             pytorch_model_config = PytorchModelConfig()
-        pytorch_model_config.setdefault("revision", "main")
+        pytorch_model_config.setdefault("revision", self.model_spec.model_revision)
         pytorch_model_config.setdefault("gpus", None)
         pytorch_model_config.setdefault("num_gpus", 1)
         pytorch_model_config.setdefault("gptq_ckpt", None)
@@ -171,7 +171,10 @@ def load(self):
             kwargs = {"torch_dtype": torch.float16}
         else:
             raise ValueError(f"Device {self._device} is not supported in temporary")
-        kwargs["revision"] = self._pytorch_model_config.get("revision", "main")
+
+        kwargs["revision"] = self._pytorch_model_config.get(
+            "revision", self.model_spec.model_revision
+        )
 
         if quantization != "none":
             if self._device == "cuda" and self._is_linux():


Test patch:
diff --git a/xinference/model/llm/pytorch/tests/test_opt.py b/xinference/model/llm/pytorch/tests/test_opt.py
--- a/xinference/model/llm/pytorch/tests/test_opt.py
+++ b/xinference/model/llm/pytorch/tests/test_opt.py
@@ -11,10 +11,14 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+import os
+from pathlib import Path
+from typing import Union
 
 import pytest
 
 from .....client import Client, GenerateModelHandle
+from ... import BUILTIN_LLM_FAMILIES
 
 
 @pytest.mark.asyncio
@@ -47,6 +51,7 @@ async def test_opt_pytorch_model(setup, quantization):
         assert isinstance(model, GenerateModelHandle)
 
         completion = model.generate("Once upon a time, there was a very old computer")
+        assert isinstance(completion, dict)
         assert "text" in completion["choices"][0]
 
         embedding_res = model.create_embedding(
@@ -56,3 +61,21 @@ async def test_opt_pytorch_model(setup, quantization):
 
         client.terminate_model(model_uid=model_uid)
         assert len(client.list_models()) == 0
+
+        # check for cached revision
+        home_address = str(Path.home())
+        snapshot_address = (
+            home_address
+            + "/.cache/huggingface/hub/models--facebook--opt-125m/snapshots"
+        )
+        actual_revision = os.listdir(snapshot_address)
+        model_name = "opt"
+        expected_revision: Union[str, None] = ""
+
+        for family in BUILTIN_LLM_FAMILIES:
+            if model_name != family.model_name:
+                continue
+            for spec in family.model_specs:
+                expected_revision = spec.model_revision
+
+        assert [expected_revision] == actual_revision
diff --git a/xinference/model/llm/tests/test_llm_family.py b/xinference/model/llm/tests/test_llm_family.py
--- a/xinference/model/llm/tests/test_llm_family.py
+++ b/xinference/model/llm/tests/test_llm_family.py
@@ -94,6 +94,7 @@ def test_serialize_llm_family_v1():
         model_size_in_billions=2,
         quantizations=["q4_0", "q4_1"],
         model_id="example/TestModel",
+        model_revision="123",
         model_file_name_template="TestModel.{quantization}.ggmlv3.bin",
     )
     pytorch_spec = PytorchLLMSpecV1(
@@ -101,6 +102,7 @@ def test_serialize_llm_family_v1():
         model_size_in_billions=3,
         quantizations=["int8", "int4", "none"],
         model_id="example/TestModel",
+        model_revision="456",
     )
     prompt_style = PromptStyleV1(
         style_name="ADD_COLON_SINGLE",
@@ -122,7 +124,7 @@ def test_serialize_llm_family_v1():
         prompt_style=prompt_style,
     )
 
-    expected = """{"version": 1, "model_name": "TestModel", "model_lang": ["en"], "model_ability": ["embed", "generate"], "model_description": null, "model_specs": [{"model_format": "ggmlv3", "model_size_in_billions": 2, "quantizations": ["q4_0", "q4_1"], "model_id": "example/TestModel", "model_file_name_template": "TestModel.{quantization}.ggmlv3.bin", "model_uri": null}, {"model_format": "pytorch", "model_size_in_billions": 3, "quantizations": ["int8", "int4", "none"], "model_id": "example/TestModel", "model_uri": null}], "prompt_style": {"style_name": "ADD_COLON_SINGLE", "system_prompt": "A chat between a curious human and an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the human's questions.", "roles": ["user", "assistant"], "intra_message_sep": "\\n### ", "inter_message_sep": "\\n### ", "stop": null, "stop_token_ids": null}}"""
+    expected = """{"version": 1, "model_name": "TestModel", "model_lang": ["en"], "model_ability": ["embed", "generate"], "model_description": null, "model_specs": [{"model_format": "ggmlv3", "model_size_in_billions": 2, "quantizations": ["q4_0", "q4_1"], "model_id": "example/TestModel", "model_revision": "123", "model_file_name_template": "TestModel.{quantization}.ggmlv3.bin", "model_uri": null}, {"model_format": "pytorch", "model_size_in_billions": 3, "quantizations": ["int8", "int4", "none"], "model_id": "example/TestModel", "model_revision": "456", "model_uri": null}], "prompt_style": {"style_name": "ADD_COLON_SINGLE", "system_prompt": "A chat between a curious human and an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the human's questions.", "roles": ["user", "assistant"], "intra_message_sep": "\\n### ", "inter_message_sep": "\\n### ", "stop": null, "stop_token_ids": null}}"""
     assert json.loads(llm_family.json()) == json.loads(expected)