Add AquilaChat2-7B-16K.py AquilaChat2-7B-16K.py.

Signed-off-by: ldwang <[email protected]>
FlagAI-Open · Oct 18, 2023 · 0178066 · 0178066
1 parent 7b610ed
commit 0178066
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 0 deletions.
diff --git a/flagai/model/aquila2_hf/AquilaChat2-34B-16K.py b/flagai/model/aquila2_hf/AquilaChat2-34B-16K.py
@@ -0,0 +1,14 @@
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import torch
+device = torch.device("cuda:0")
+model_info = "BAAI/AquilaChat2-34B-16k"
+tokenizer = AutoTokenizer.from_pretrained(model_info, trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained(model_info, trust_remote_code=True, torch_dtype=torch.bfloat16)
+model.eval()
+model.to(device)
+text = "请给出10个要到北京旅游的理由。"
+from predict import predict
+out = predict(model, text, tokenizer=tokenizer, max_gen_len=200, top_p=0.95,
+              seed=1234, topk=100, temperature=0.9, sft=True, device=device,
+              model_name="AquilaChat2-34B-16K")
+print(out)
diff --git a/flagai/model/aquila2_hf/AquilaChat2-7B-16K.py b/flagai/model/aquila2_hf/AquilaChat2-7B-16K.py
@@ -0,0 +1,14 @@
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import torch
+device = torch.device("cuda:0")
+model_info = "BAAI/AquilaChat2-7B-16K"
+tokenizer = AutoTokenizer.from_pretrained(model_info, trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained(model_info, trust_remote_code=True, torch_dtype=torch.bfloat16)
+model.eval()
+model.to(device)
+text = "请给出10个要到北京旅游的理由。"
+from predict import predict
+out = predict(model, text, tokenizer=tokenizer, max_gen_len=200, top_p=0.95,
+              seed=1234, topk=100, temperature=0.9, sft=True, device=device,
+              model_name="AquilaChat2-7B-16K")
+print(out)