hjxdl 0.0.1__tar.gz → 0.0.2__tar.gz

{hjxdl-0.0.1/hjxdl.egg-info → hjxdl-0.0.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: hjxdl
-Version: 0.0.1
+Version: 0.0.2
 Summary: A collection of functions for Jupyter notebooks
 Home-page: https://github.com/huluxiaohuowa/hdl
 Author: Jianxing Hu

{hjxdl-0.0.1 → hjxdl-0.0.2}/hdl/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.0.1'
-__version_tuple__ = version_tuple = (0, 0, 1)
+__version__ = version = '0.0.2'
+__version_tuple__ = version_tuple = (0, 0, 2)

hjxdl-0.0.2/hdl/utils/llm/chat.py

@@ -0,0 +1,164 @@
+import typing as t
+
+from llama_cpp import Llama
+from jupyfuncs.llm.openapi import (
+    chat_oai_invoke,
+    chat_oai_stream
+)
+
+
+class GGUF_M():
+    def __init__(
+        self,
+        model_path :str,
+        device: str='gpu',
+        generation_kwargs: dict = {},
+        server_ip: str = "127.0.0.1",
+        server_port: int = 8000,
+    ):
+        """Initialize the model from a local path.
+        Here it does not manage chat histories.
+        
+        Args:
+            model_path (str): The path to the model.
+            device (str, optional): The device to use for model initialization. Defaults to 'gpu'.
+            generation_kwargs (dict, optional): Additional keyword arguments for model generation. Defaults to {}.
+        """
+        # 从本地初始化模型
+        super().__init__()
+        self.generation_kwargs = generation_kwargs
+        self.questions = []
+        self.resps = []
+        print("正在从本地加载模型...")
+        if device == 'cpu':
+            self.model = Llama(
+                model_path=model_path,
+                n_threads=self.generation_kwargs['num_threads'],
+                n_ctx=self.generation_kwargs['max_context_length'],
+            )
+        else:
+            self.model = Llama(
+                model_path=model_path,
+                n_threads=self.generation_kwargs['num_threads'],
+                n_ctx=self.generation_kwargs['max_context_length'],
+                n_gpu_layers=-1
+            )
+            
+        print("完成本地模型的加载")
+
+    def invoke(
+        self, 
+        prompt : str,
+        stop: list[str] | None = ["USER:", "ASSISTANT:"],
+        # history: list = [],
+        **kwargs: t.Any,
+    ) -> str:
+        """Invoke the model to generate a response based on the given prompt.
+        
+        Args:
+            prompt (str): The prompt to be used for generating the response.
+            stop (list[str], optional): List of strings that indicate when the model should stop generating the response. Defaults to ["USER:", "ASSISTANT:"].
+            **kwargs: Additional keyword arguments to be passed to the model.
+        
+        Returns:
+            str: The generated response based on the prompt.
+        """
+        prompt_final = f"USER:\n{prompt}\nASSISTANT:\n" 
+
+        result = self.model.create_completion(
+            prompt_final,
+            repeat_penalty=self.generation_kwargs["repetition_penalty"],
+            max_tokens=self.generation_kwargs["max_new_tokens"],
+            stop=stop,
+            echo=False,
+            temperature=self.generation_kwargs["temperature"],
+            mirostat_mode = 2,
+            mirostat_tau=4.0,
+            mirostat_eta=1.1
+        )
+        resp = result['choices'][0]['text']
+        # history.append(
+        #     [prompt, resp]
+        # )
+        return resp
+ 
+    def stream(
+        self,
+        prompt: str,
+        stop: list[str] | None = ["USER:", "ASSISTANT:"],
+        # history: list = [],
+        **kwargs: t.Any,
+    ):
+        """Generate text responses based on the given prompt using the model.
+        
+        Args:
+            prompt (str): The prompt to generate text responses.
+            stop (list[str], optional): List of strings to stop the generation. Defaults to ["USER:", "ASSISTANT:"].
+            **kwargs: Additional keyword arguments for the model.
+        
+        Yields:
+            str: Text responses generated by the model based on the prompt.
+        """
+        self.questions.append(prompt)
+        prompt = f"USER:\n{prompt}\nASSISTANT:\n" 
+        output = self.model.create_completion(
+            prompt,
+            stream=True,
+            repeat_penalty=self.generation_kwargs["repetition_penalty"], 
+            max_tokens=self.generation_kwargs["max_new_tokens"],
+            stop=stop,
+            echo=False,
+            temperature=self.generation_kwargs["temperature"],
+            mirostat_mode = 2,
+            mirostat_tau=4.0,
+            mirostat_eta=1.1
+        )
+        # history.append([])
+        for chunk in output:
+            item = chunk['choices'][0]['text']
+            # self.resps[-1].append(item)
+            yield chunk['choices'][0]['text']
+        # self.resps[-1] = "".join(self.resps[-1])
+
+
+class OpenAI_M():
+    def __init__(
+        self,
+        model_path: str = None,
+        device: str='gpu',
+        generation_kwargs: dict = {},
+        server_ip: str = "172.28.1.2",
+        server_port: int = 8000,
+    ):
+        self.model_path = model_path
+        self.server_ip = server_ip
+        self.server_port = server_port
+        self.base_url = "http://{self.server_ip}:{str(self.server_port)}/v1"
+    
+    def invoke(
+        self, 
+        prompt : str,
+        stop: list[str] | None = ["USER:", "ASSISTANT:"],
+        # history: list = [],
+        **kwargs: t.Any,
+    ) -> str:
+        resp = chat_oai_invoke(
+            base_url=self.base_url,
+            model=self.model_path,
+            prompt=prompt
+        )
+        return resp
+    
+    def stream(
+        self, 
+        prompt : str,
+        stop: list[str] | None = ["USER:", "ASSISTANT:"],
+        # history: list = [],
+        **kwargs: t.Any,
+    ) -> str:
+        resp = chat_oai_stream(
+            base_url=self.base_url,
+            model=self.model_path,
+            prompt=prompt
+        )
+        return resp

{hjxdl-0.0.1 → hjxdl-0.0.2/hjxdl.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: hjxdl
-Version: 0.0.1
+Version: 0.0.2
 Summary: A collection of functions for Jupyter notebooks
 Home-page: https://github.com/huluxiaohuowa/hdl
 Author: Jianxing Hu

{hjxdl-0.0.1 → hjxdl-0.0.2}/hjxdl.egg-info/SOURCES.txt

@@ -102,6 +102,7 @@ hdl/utils/database_tools/__init__.py
 hdl/utils/database_tools/connect.py
 hdl/utils/general/__init__.py
 hdl/utils/general/glob.py
+hdl/utils/llm/chat.py
 hdl/utils/schedulers/__init__.py
 hdl/utils/schedulers/norm_lr.py
 hjxdl.egg-info/PKG-INFO

hjxdl-0.0.2/version.txt

@@ -0,0 +1 @@
+0.0.2

hjxdl-0.0.1/version.txt

@@ -1 +0,0 @@
-0.0.1