release v0.2.4

Author: hiyouga

requirements.txt

@@ -1,4 +1,7 @@
-vllm>=0.4.2
-infinity-emb[torch,optimum]>=0.0.33
+vllm==0.4.3
+infinity-emb[torch,optimum]==0.0.40
 openai>=1.5.0
+uvicorn
+pydantic
+fastapi
 sse-starlette

src/imitater/__init__.py

@@ -1 +1 @@
-__version__ = "0.2.4.dev0"
+__version__ = "0.2.4"

src/imitater/model/chat_model.py

@@ -34,14 +34,14 @@ class ChatConfig:
 
     @staticmethod
     def add_cli_args(parser: "ArgumentParser") -> None:
-        parser.add_argument("--name", type=str)
-        parser.add_argument("--path", type=str)
-        parser.add_argument("--device", type=int, nargs="+")
-        parser.add_argument("--port", type=int)
-        parser.add_argument("--maxlen", type=int, default=2048)
-        parser.add_argument("--agent_type", type=str, choices=list_agents(), default="react")
-        parser.add_argument("--template", type=str, default=None)
-        parser.add_argument("--gen_config", type=str, default=None)
+        parser.add_argument("--name", type=str, required=True, help="Model name.")
+        parser.add_argument("--path", type=str, required=True, help="Model path or hub id.")
+        parser.add_argument("--device", type=int, nargs="+", required=True, help="Device ids.")
+        parser.add_argument("--port", type=int, required=True, help="API port.")
+        parser.add_argument("--maxlen", type=int, default=2048, help="Max sequence length.")
+        parser.add_argument("--agent_type", type=str, choices=list_agents(), default="react", help="Agent type.")
+        parser.add_argument("--template", type=str, default=None, help="Template path.")
+        parser.add_argument("--gen_config", type=str, default=None, help="Generation config path.")
 
     @classmethod
     def from_cli_args(cls, args: "Namespace") -> Self:
@@ -130,7 +130,7 @@ async def _generate(self, messages: List[Dict[str, str]], request_id: str, **gen
             stop_token_ids=self._generation_config.eos_token_id + gen_kwargs.pop("stop_token_ids", []),
         )
         result_generator = self._engine.generate(
-            prompt=None, sampling_params=sampling_params, request_id=request_id, prompt_token_ids=input_ids
+            inputs={"prompt_token_ids": input_ids}, sampling_params=sampling_params, request_id=request_id
         )
         return result_generator
 

src/imitater/model/embed_model.py

@@ -32,11 +32,11 @@ class EmbedConfig:
 
     @staticmethod
     def add_cli_args(parser: "ArgumentParser") -> None:
-        parser.add_argument("--name", type=str, required=True)
-        parser.add_argument("--path", type=str, required=True)
-        parser.add_argument("--device", type=int, nargs="+", required=True)
-        parser.add_argument("--port", type=int, required=True)
-        parser.add_argument("--batch_size", type=int, default=64)
+        parser.add_argument("--name", type=str, required=True, help="Model name.")
+        parser.add_argument("--path", type=str, required=True, help="Model path or hub id.")
+        parser.add_argument("--device", type=int, nargs="+", required=True, help="Device ids.")
+        parser.add_argument("--port", type=int, required=True, help="API port.")
+        parser.add_argument("--batch_size", type=int, default=64, help="Batch size.")
 
     @classmethod
     def from_cli_args(cls, args: "Namespace") -> Self: