Initial commit

scripts/chat.py

@@ -0,0 +1,32 @@
+import os
+import torch
+from khaosz import Khaosz
+
+
+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)
+
+    histroy = []
+    while True:
+        query = input(">> ")
+        if query == "!exit":
+            break
+        
+        response_size = 0
+        for response, histroy in model.stream_generate(
+            query=query, 
+            history=histroy,
+            temperature=0.7,
+            top_p=0.95,
+            top_k=30
+        ):
+            print(response[response_size:], end="", flush=True)
+            response_size = len(response)
+    
+
+if __name__ == "__main__":
+    chat()

scripts/download.py

@@ -0,0 +1,14 @@
+import os
+from huggingface_hub import snapshot_download
+
+
+PROJECT_ROOT = os.path.dirname(
+    os.path.dirname(os.path.abspath(__file__)))
+
+
+if __name__ == "__main__":
+    snapshot_download(
+        repo_id="ViperEk/KHAOSZ",
+        local_dir=os.path.join(PROJECT_ROOT, "params"),  
+        force_download=True
+    )

scripts/generate_ar.py

@@ -0,0 +1,27 @@
+import os
+import torch
+from khaosz import Khaosz
+
+
+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)
+
+    query = input(">> ")
+    
+    response = model.text_generate(
+        query=query, 
+        temperature=0.6,
+        top_p=0.95,
+        top_k=30
+    )
+    
+    print(response)
+
+    
+
+if __name__ == "__main__":
+    generate_text()

scripts/generate_batch.py

@@ -0,0 +1,25 @@
+import os
+import torch
+from khaosz import Khaosz
+
+
+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)
+    inputs = ["你好", "请问什么是人工智能", "今天天气如何", "我感到焦虑， 请问我应该怎么办", "请问什么是显卡"]
+    
+    responses = model.batch_generate(
+        queries=inputs,
+        temperature=0.7,
+        top_p=0.95,
+        top_k=30
+    )
+    
+    for q, r in zip(inputs, responses):
+        print((q, r))
+
+if __name__ == "__main__":
+    batch_generate()

scripts/generate_retrieve.py

@@ -0,0 +1,42 @@
+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)
+    
+    retrive_response = model.retrieve_generate(
+        retrieved=retrieved,
+        query=query,
+        temperature=0.7,
+        top_k=30,
+        top_p=0.95,
+    )
+
+    print("retrive content:")
+    print("\n".join([f"{idx + 1}. " + text for idx, (text, _) in enumerate(retrieved)]))
+
+    print("\n\nretrive generate:")
+    print(retrive_response)