Update app.py

This commit is contained in:
Baize 2023-04-05 01:43:02 +00:00 committed by huggingface-web
parent db8c43fa76
commit 95b9be1079
1 changed files with 2 additions and 0 deletions

2
app.py
View File

@ -47,6 +47,8 @@ def predict(text,
global total_count
total_count += 1
print(total_count)
if total_count % 50 == 0 :
os.system("nvidia-smi")
with torch.no_grad():
for x in greedy_search(input_ids,model,tokenizer,stop_words=["[|Human|]", "[|AI|]"],max_length=max_length_tokens,temperature=temperature,top_p=top_p):
if is_stop_word_or_prefix(x,["[|Human|]", "[|AI|]"]) is False: