wmpscc commited on
Commit
8596d50
·
1 Parent(s): 7e84f81

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +2 -2
app.py CHANGED
@@ -8,7 +8,7 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
8
 
9
  def init_model():
10
  model = AutoModelForCausalLM.from_pretrained("Linly-AI/Chinese-LLaMA-2-7B-hf", device_map="cuda:0", torch_dtype=torch.float16, trust_remote_code=True)
11
- tokenizer = AutoTokenizer.from_pretrained("Linly-AI/Chinese-LLaMA-2-7B-hf", use_fast=False, trust_remote_code=True)
12
  return model, tokenizer
13
 
14
 
@@ -16,7 +16,7 @@ def init_model():
16
  def chat(prompt, top_k, temperature):
17
  prompt = f"### Instruction:{prompt.strip()} ### Response:"
18
  inputs = tokenizer(prompt, return_tensors="pt").to("cuda:0")
19
- generate_ids = model.generate(inputs.input_ids, max_new_tokens=2048, do_sample = True, top_k=int(top_k), top_p=0, temperature=float(temperature), repetition_penalty=1.15, eos_token_id=2, bos_token_id=1, pad_token_id=0)
20
  response = tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
21
  response = response.lstrip(prompt)
22
  print('-log: ',prompt, response)
 
8
 
9
  def init_model():
10
  model = AutoModelForCausalLM.from_pretrained("Linly-AI/Chinese-LLaMA-2-7B-hf", device_map="cuda:0", torch_dtype=torch.float16, trust_remote_code=True)
11
+ tokenizer = AutoTokenizer.from_pretrained("Linly-AI/Chinese-LLaMA-2-7B-hf", use_fast=True, trust_remote_code=True)
12
  return model, tokenizer
13
 
14
 
 
16
  def chat(prompt, top_k, temperature):
17
  prompt = f"### Instruction:{prompt.strip()} ### Response:"
18
  inputs = tokenizer(prompt, return_tensors="pt").to("cuda:0")
19
+ generate_ids = model.generate(inputs.input_ids, do_sample=True, max_new_tokens=2048, top_k=int(top_k), top_p=0.84, temperature=float(temperature), repetition_penalty=1.15, eos_token_id=2, bos_token_id=1, pad_token_id=0)
20
  response = tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
21
  response = response.lstrip(prompt)
22
  print('-log: ',prompt, response)