can't stop generation words

Nobody has claimed this yet.

I wrote a chat api to inference based on the V0.2.5. and downloaded the weight v1.1. but when I run, I asked a simple question
hello, tell me your name or 1+1, it outputs many words that result in OOM.I see issue433 solution, but it does't work.here is my code,

import time

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, LlamaTokenizer, AutoModel,AutoModelForSeq2SeqLM
import docx
import os
import argparse
import pandas as pd
import gc

def load_model(
model_path, device, num_gpus, max_gpu_memory=None, load_8bit=False, debug=False
):
if device == "cpu":
kwargs = {}
elif device == "cuda":
kwargs = {"torch_dtype": torch.float16}
if num_gpus == "auto":
kwargs["device_map"] = "auto"
else:
num_gpus = int(num_gpus)
if num_gpus != 1:
kwargs["device_map"] = "auto"
if max_gpu_memory is None:
kwargs[
"device_map"
] = "sequential" # This is important for not the same VRAM sizes
available_gpu_memory = get_gpu_memory(num_gpus)
kwargs["max_memory"] = {
i: str(int(available_gpu_memory[i] * 0.85)) + "GiB"
for i in range(num_gpus)
}
else:
kwargs["max_memory"] = {i: max_gpu_memory for i in range(num_gpus)}
print("init_kwargs", kwargs)
else:
raise ValueError(f"Invalid device: {device}")

tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=False)
model = AutoModelForCausalLM.from_pretrained(
model_path, low_cpu_mem_usage=True, **kwargs
)
# raise_warning_for_old_weights(model_path, model)

if (device == "cuda" and num_gpus == 1) or device == "mps":
model.to(device)

if debug:
print(model)

return model, tokenizer

def generate_stream(
model, tokenizer, params, device, context_len=2048,stream_interval=2):
prompt = params["prompt"]
len_prompt = len(prompt)
temperature = float(params.get("temperature", 1.0))
max_new_tokens = int(params.get("max_new_tokens", 256))
stop_str = params.get("stop", None)
echo = params.get("echo", True)
stop_token_ids = params.get("stop_token_ids", None) or []
stop_token_ids.append(tokenizer.eos_token_id)

input_ids = tokenizer(prompt).input_ids
input_echo_len = len(input_ids)
output_ids = list(input_ids)

if model.config.is_encoder_decoder:
max_src_len = context_len
else:
max_src_len = context_len - max_new_tokens - 8

input_ids = input_ids[-max_src_len:]

if model.config.is_encoder_decoder:
encoder_output = model.encoder(input_ids=torch.as_tensor([input_ids],
device=device))[0]
start_ids = torch.as_tensor([[model.generation_config.decoder_start_token_id]],
dtype=torch.int64, device=device)

for i in range(max_new_tokens):
if i == 0:
if model.config.is_encoder_decoder:
out = model.decoder(input_ids=start_ids,
encoder_hidden_states=encoder_output,
use_cache=True)
logits = model.lm_head(out[0])
else:
out = model(torch.as_tensor([input_ids], device=device), use_cache=True)
logits = out.logits
past_key_values = out.past_key_values
else:
if model.config.is_encoder_decoder:
out = model.decoder(input_ids=torch.as_tensor([[token]], device=device),
encoder_hidden_states=encoder_output,
use_cache=True,
past_key_values=past_key_values)

logits = model.lm_head(out[0])
else:
out = model(
input_ids=torch.as_tensor([[token]], device=device),
use_cache=True,
past_key_values=past_key_values,
)
logits = out.logits
past_key_values = out.past_key_values

last_token_logits = logits[0][-1]

if device == "mps":
# Switch to CPU by avoiding some bugs in mps backend.
last_token_logits = last_token_logits.float().to("cpu")

if temperature < 1e-4:
token = int(torch.argmax(last_token_logits))
else:
probs = torch.softmax(last_token_logits / temperature, dim=-1)
token = int(torch.multinomial(probs, num_samples=1))

output_ids.append(token)

if token in stop_token_ids:
stopped = True
else:
stopped = False

if i % stream_interval == 0 or i == max_new_tokens - 1 or stopped:
if echo:
tmp_output_ids = output_ids
rfind_start = len_prompt
else:
tmp_output_ids = output_ids[input_echo_len:]
rfind_start = 0

output = tokenizer.decode(tmp_output_ids, skip_special_tokens=True,
spaces_between_special_tokens=False)
if stop_str:
pos = output.rfind(stop_str, rfind_start)
if pos != -1:
output = output[:pos]
stopped = True
yield output

if stopped:
break

del past_key_values, out
gc.collect()
torch.cuda.empty_cache()

class ChatAPI(object):
def __init__(self,model_path="/workspace/fashchat_model_v1.1/",temperature=0.7,max_input_tokens=2048,max_output_token=512):
self.model_path = model_path
self.temperature = temperature
self.max_input_tokens = max_input_tokens
self.max_output_token = max_output_token
self.model, self.tokenizer = load_model(
self.model_path, "cuda", 1
)
self.params = {
"temperature": self.temperature, # 模型输出的随机程度
"max_new_tokens": self.max_output_token, # AI最大回复的语言长度
"stop": None,
'stop_token_ids': None,
'echo': False
}

def start_chat(self,prompt):
question = {"prompt":prompt}
self.params.update(question)
iter = generate_stream(self.model, self.tokenizer, self.params, 'cuda', context_len=self.max_input_tokens)

pre = 0
for outputs in iter:
outputs = outputs.strip().split(" ")
now = len(outputs) - 1
if now > pre:
print(" ".join(outputs[pre:now]), end=" ", flush=True)
pre = now
print(" ".join(outputs[pre:]), flush=True)
return " ".join(outputs)

if __name__ == '__main__':
chat_api = ChatAPI()
result = chat_api.start_chat("1+1")
print("result",result.strip())

I guess the problem may be the stop words, but I don't know how to solve it.

Start at the provided generate_stream function, especially its max_new_tokens and stop_token_ids handling, then trace how ChatAPI.start_chat sets those parameters. Reproduce the unbounded-looking response with the shown 1+1 prompt and determine whether generation stops at an EOS or configured stop condition without exhausting GPU memory.

Written by the indexing model from the issue text.