added local llm class.

Author: JDOxygen

examples/llms/demo_local_llm.py

@@ -0,0 +1,23 @@
+import asyncio
+
+from oxygent import MAS, oxy
+
+oxy_space = [
+    oxy.LocalLLM(
+        name="default_llm",
+        model_path="/path/to/your_model",
+    ),
+    oxy.ChatAgent(
+        name="master_agent",
+        llm_model="default_llm",
+    ),
+]
+
+
+async def main():
+    async with MAS(oxy_space=oxy_space) as mas:
+        await mas.start_web_service(first_query="hello")
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

oxygent/oxy/llms/local_llm.py

@@ -11,8 +11,8 @@
 
 class LocalLLM(BaseLLM):
     model_path: str = Field("")
-    device: str = Field("auto")
-    model_name: str = Field("")
+    device_map: str = Field("auto")
+    dtype: str = Field("bfloat16")
 
     async def init(self):
         try:
@@ -27,20 +27,28 @@ async def init(self):
         await super().init()
         # Load model directly
         self._model = AutoModelForCausalLM.from_pretrained(
-            self.model_path, device_map=self.device, torch_dtype=torch.bfloat16
+            self.model_path, device_map=self.device_map, dtype=self.dtype
         )
         self._tokenizer = AutoTokenizer.from_pretrained(self.model_path)
 
     async def _execute(self, oxy_request: OxyRequest) -> OxyResponse:
-        payload = {"model": self.model_name, "stream": False}
-        payload.update(Config.get_llm_config())
+        payload = Config.get_llm_config()
         for k, v in self.llm_params.items():
             payload[k] = v
         for k, v in oxy_request.arguments.items():
             if k == "messages":
                 continue
             payload[k] = v
-        payload = {"max_new_tokens": 512}
+
+        replace_dict = {
+            "max_tokens": "max_new_tokens",
+            "stream": "",
+        }
+        for k, v in replace_dict.items():
+            if k in payload:
+                if v:
+                    payload[v] = payload[k]
+                del payload[k]
 
         messages = oxy_request.arguments["messages"]