Function chat

SwissArmyTransformer/examples/chatglm3/chat_sat.py:31–52

(model, tokenizer, 
        max_length: int = 256, num_beams=1, top_p=0.7, top_k=0, temperature=0.95)

    return content, history

def chat(model, tokenizer, 
        max_length: int = 256, num_beams=1, top_p=0.7, top_k=0, temperature=0.95):
    query = "你好"
    history = []
    inputs = tokenizer.build_chat_input(query, history=history, role="user")['input_ids'].to(model.parameters().__next__().device)[0]
    seq = torch.cat(
        [inputs, torch.tensor([-1]*(max_length-len(inputs)), device=inputs.device)], dim=0
    )
    strategy = BaseStrategy(temperature=temperature, top_p=top_p, top_k=0, end_tokens=[tokenizer.eos_token_id, tokenizer.get_command("<|user|>"), tokenizer.get_command("<|observation|>")])

    output = filling_sequence(
        model, seq,
        batch_size=1,
        strategy=strategy
    )[0] # drop memory
    output_list = output.tolist()[0][len(inputs):-1]

    response = tokenizer.decode(output_list)
    history.append({"role": "user", "content": query})
    response, history = process_response(response, history)
    print(response)
    print(history)

if __name__ == "__main__":
    import argparse