refactor: 优化工具脚本接口并修复批处理问题
This commit is contained in:
@@ -13,11 +13,9 @@ PARAMETER_ROOT = Path(PROJECT_ROOT, "params")
|
||||
def generate_text():
|
||||
# Load model from pretrained
|
||||
model = AutoModel.from_pretrained(PARAMETER_ROOT)
|
||||
tokenizer = AutoTokenizer.from_pretrained(PARAMETER_ROOT)
|
||||
model.to(device="cuda", dtype=torch.bfloat16)
|
||||
|
||||
# Load tokenizer from pretrained
|
||||
tokenizer = AutoTokenizer.from_pretrained(PARAMETER_ROOT / "tokenizer")
|
||||
|
||||
query = input(">> ")
|
||||
|
||||
engine = InferenceEngine(
|
||||
|
||||
Reference in New Issue
Block a user