chore: 修改文件夹结构

scripts/demo/download.py

@@ -0,0 +1,12 @@
+from pathlib import Path
+from huggingface_hub import snapshot_download
+
+PROJECT_ROOT = Path(__file__).parent.parent
+PARAMETER_ROOT = Path(PROJECT_ROOT, "params")
+
+if __name__ == "__main__":
+    snapshot_download(
+        repo_id="ViperEk/AstrAI",
+        local_dir=PARAMETER_ROOT,
+        force_download=True,
+    )

scripts/demo/generate_ar.py

@@ -0,0 +1,35 @@
+import torch
+from pathlib import Path
+from astrai.config.param_config import ModelParameter
+from astrai.inference.core import disable_random_init
+from astrai.inference.generator import GeneratorFactory, GenerationRequest
+
+PROJECT_ROOT = Path(__file__).parent.parent
+PARAMETER_ROOT = Path(PROJECT_ROOT, "params")
+
+
+def generate_text():
+
+    with disable_random_init():
+        param = ModelParameter.load(PARAMETER_ROOT)
+        param.to(device="cuda", dtype=torch.bfloat16)
+
+    query = input(">> ")
+
+    request = GenerationRequest(
+        query=query,
+        temperature=0.8,
+        top_p=0.95,
+        top_k=50,
+        max_len=param.config.max_len,
+        history=None,
+        system_prompt=None,
+    )
+    generator = GeneratorFactory.create(param, request)
+    response = generator.generate(request)
+
+    print(response)
+
+
+if __name__ == "__main__":
+    generate_text()

scripts/demo/generate_batch.py

@@ -0,0 +1,42 @@
+import torch
+from pathlib import Path
+from astrai.config.param_config import ModelParameter
+from astrai.inference.core import disable_random_init
+from astrai.inference.generator import GeneratorFactory, GenerationRequest
+
+PROJECT_ROOT = Path(__file__).parent.parent
+PARAMETER_ROOT = Path(PROJECT_ROOT, "params")
+
+
+def batch_generate():
+
+    with disable_random_init():
+        param = ModelParameter.load(PARAMETER_ROOT)
+        param.to(device="cuda", dtype=torch.bfloat16)
+
+    inputs = [
+        "你好",
+        "请问什么是人工智能",
+        "今天天气如何",
+        "我感到焦虑， 请问我应该怎么办",
+        "请问什么是显卡",
+    ]
+
+    request = GenerationRequest(
+        query=inputs,
+        temperature=0.8,
+        top_p=0.95,
+        top_k=50,
+        max_len=param.config.max_len,
+        history=None,
+        system_prompt=None,
+    )
+    generator = GeneratorFactory.create(param, request)
+    responses = generator.generate(request)
+
+    for q, r in zip(inputs, responses):
+        print((q, r))
+
+
+if __name__ == "__main__":
+    batch_generate()

scripts/demo/stream_chat.py

@@ -0,0 +1,47 @@
+import torch
+from pathlib import Path
+from astrai.config.param_config import ModelParameter
+from astrai.inference.core import disable_random_init
+from astrai.inference.generator import GeneratorFactory, GenerationRequest
+
+PROJECT_ROOT = Path(__file__).parent.parent
+PARAMETER_ROOT = Path(PROJECT_ROOT, "params")
+
+
+def chat():
+
+    with disable_random_init():
+        param = ModelParameter.load(PARAMETER_ROOT)
+        param.to(device="cuda", dtype=torch.bfloat16)
+
+    history = []
+    while True:
+        query = input(">> ")
+        if query == "!exit":
+            break
+
+        request = GenerationRequest(
+            query=query,
+            temperature=0.8,
+            top_p=0.95,
+            top_k=50,
+            max_len=param.config.max_len,
+            history=history,
+            system_prompt=None,
+        )
+        generator = GeneratorFactory.create(param, request)
+
+        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()