Update app.py

app.py

@@ -7,7 +7,7 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
 
 
 def init_model():
-    model = AutoModelForCausalLM.from_pretrained("Linly-AI/Chinese-LLaMA-2-7B-hf", device_map="auto", torch_dtype=torch.float16, trust_remote_code=True)
+    model = AutoModelForCausalLM.from_pretrained("Linly-AI/Chinese-LLaMA-2-7B-hf", device_map="cuda:0", torch_dtype=torch.float16, trust_remote_code=True)
     tokenizer = AutoTokenizer.from_pretrained("Linly-AI/Chinese-LLaMA-2-7B-hf", use_fast=False, trust_remote_code=True)
     return model, tokenizer
     
@@ -15,8 +15,8 @@ def init_model():
 
 def chat(prompt, top_k, temperature):
     prompt = f"### Instruction:{prompt.strip()}  ### Response:"
-    inputs = tokenizer(prompt, return_tensors="pt")
-    generate_ids = model.generate(inputs.input_ids, max_new_tokens=2048, do_sample = True, top_k=top_k, top_p = 0, temperature=temperature, repetition_penalty=1.15, eos_token_id=2, bos_token_id=1, pad_token_id=0)
+    inputs = tokenizer(prompt, return_tensors="pt").to("cuda:0")
+    generate_ids = model.generate(inputs.input_ids, max_new_tokens=20480, do_sample = True, top_k=top_k, top_p = 0, temperature=temperature, repetition_penalty=1.15, eos_token_id=2, bos_token_id=1, pad_token_id=0)
     response = tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
     response = response.lstrip(prompt)
     print('log:', response)