feat: 增加推理部分工厂模式

demo/generate_ar.py

@@ -1,27 +1,35 @@
 import os
 import torch
-from khaosz import Khaosz
+from khaosz.config.param_config import ModelParameter
+from khaosz.inference.core import disable_random_init
+from khaosz.inference.generator import LoopGenerator, GenerationRequest
 
 
 PROJECT_ROOT = os.path.dirname(
     os.path.dirname(os.path.abspath(__file__)))
 
 def generate_text():
-    model_dir = os.path.join(PROJECT_ROOT, "params")
-    model = Khaosz(model_dir).to(device='cuda', dtype=torch.bfloat16)
-
+    
+    with disable_random_init():
+        model_dir = os.path.join(PROJECT_ROOT, "params")
+        param = ModelParameter.load(model_dir)
+    
+    param.to(device='cuda', dtype=torch.bfloat16)
     query = input(">> ")
     
-    response = model.text_generate(
-        query=query, 
+    request = GenerationRequest(
+        query=query,
         temperature=0.8,
         top_p=0.95,
-        top_k=50
+        top_k=50,
+        max_len=param.config.max_len,
+        history=None,
+        system_prompt=None,
     )
+    generator = LoopGenerator(param)
+    response = generator.generate(request)
     
     print(response)
 
-    
-
 if __name__ == "__main__":
     generate_text()

demo/generate_batch.py

@@ -1,22 +1,31 @@
 import os
 import torch
-from khaosz import Khaosz
-
+from khaosz.config.param_config import ModelParameter
+from khaosz.inference.core import disable_random_init
+from khaosz.inference.generator import BatchGenerator, GenerationRequest
 
 PROJECT_ROOT = os.path.dirname(
     os.path.dirname(os.path.abspath(__file__)))
 
 def batch_generate():
-    model_dir = os.path.join(PROJECT_ROOT, "params")
-    model = Khaosz(model_dir).to(device='cuda', dtype=torch.bfloat16)
+    with disable_random_init():
+        model_dir = os.path.join(PROJECT_ROOT, "params")
+        param = ModelParameter.load(model_dir)
+
+    param.to(device='cuda', dtype=torch.bfloat16)
+    generator = BatchGenerator(param)
     inputs = ["你好", "请问什么是人工智能", "今天天气如何", "我感到焦虑， 请问我应该怎么办", "请问什么是显卡"]
     
-    responses = model.batch_generate(
+    request = GenerationRequest(
         query=inputs,
         temperature=0.8,
         top_p=0.95,
-        top_k=50
+        top_k=50,
+        max_len=param.config.max_len,
+        history=None,
+        system_prompt=None,
     )
+    responses = generator.generate(request)
     
     for q, r in zip(inputs, responses):
         print((q, r))

demo/generate_retrieve.py

@@ -1,43 +0,0 @@
-import os
-import torch
-from khaosz import Khaosz, SemanticTextSplitter, Retriever
-
-
-PROJECT_ROOT = os.path.dirname(
-    os.path.dirname(os.path.abspath(__file__)))
-
-if __name__ == "__main__":
-    model_dir = os.path.join(PROJECT_ROOT, "params")
-    context_path = os.path.join(PROJECT_ROOT, "README.md")
-    
-    model = Khaosz(model_dir).to(device='cuda', dtype=torch.bfloat16)
-    spliter = SemanticTextSplitter(model.encode)
-    retriever = Retriever()
-    text = open(context_path, "r", encoding="utf-8").read()
-    
-    res = spliter.split(text, threshold=0.8, window_size=1)
-    # print(("\n" + "+"*100 + "\n").join(res))
-
-    res_embs = model.encode(res)
-    for sentence, emb in zip(res, res_embs):
-        retriever.add_vector(sentence, emb)
-
-    retrive_top_k = 5
-    query = "作者设计了一个怎样的模型"
-    emb_query = model.encode(query)
-    retrieved = retriever.retrieve(emb_query, retrive_top_k)
-    retrieved_content = "\n".join([f"{idx + 1}. " + text for idx, (text, _) in enumerate(retrieved)])
-    
-    retrive_response = model.retrieve_generate(
-        retrieved=retrieved_content,
-        query=query,
-        temperature=0.8,
-        top_p=0.95,
-        top_k=50
-    )
-
-    print("retrieve content:")
-    print(retrieved_content)
-
-    print("\n\nretrive generate:")
-    print(retrive_response)

demo/stream_chat.py

@@ -1,14 +1,21 @@
 import os
 import torch
-from khaosz import Khaosz
+from khaosz.config.param_config import ModelParameter
+from khaosz.inference.core import disable_random_init
+from khaosz.inference.generator import StreamGenerator, GenerationRequest
 
 
 PROJECT_ROOT = os.path.dirname(
     os.path.dirname(os.path.abspath(__file__)))
 
 def chat():
-    model_dir = os.path.join(PROJECT_ROOT, "params")
-    model = Khaosz(model_dir).to(device='cuda', dtype=torch.bfloat16)
+    
+    with disable_random_init():
+        model_dir = os.path.join(PROJECT_ROOT, "params")
+        param = ModelParameter.load(model_dir)
+
+    param.to(device='cuda', dtype=torch.bfloat16)
+    generator = StreamGenerator(param)
 
     history = []
     while True:
@@ -16,17 +23,27 @@ def chat():
         if query == "!exit":
             break
         
-        response_size = 0
-        for response, history in model.stream_generate(
-            query=query, 
-            history=history,
+        request = GenerationRequest(
+            query=query,
             temperature=0.8,
             top_p=0.95,
-            top_k=50
-        ):
+            top_k=50,
+            max_len=param.config.max_len,
+            history=history,
+            system_prompt=None,
+        )
+        
+        response_size = 0
+        full_response = ""
+        for response in generator.generate(request):
+            # response is the cumulative response up to current token
             print(response[response_size:], end="", flush=True)
             response_size = len(response)
-    
+            full_response = response
+        
+        # After generation, update history
+        history.append((query, full_response.strip()))
+
 
 if __name__ == "__main__":
     chat()