chore: 修改文件夹结构

2026-03-31 10:14:08 +08:00
parent b1527d9575
commit 4ead0a20cf
10 changed files with 14 additions and 20 deletions
@@ -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()