allenai

/

MolmoE-1B-0924

like 158

Follow

Ai2 6.03k

Hey Team,

I'm finding hard to properly integrate chat history as shown in the Molmo Chat WebPage.

Would be great if you could help me out with this, :>)

chat_history = []

while True:
    print(f"{color_b}************ Enter Your Input Query: *****************{reset}")
    input_query = input()
    query_text = input_query + "\n"+ f"{chat_history}" if chat_history else input_query


    print("######################")

    if input_query == exit_trigger:
        print(f"{color_r}**************************Exiting Molmo Chat!***************************{reset}")
        break


    inputs = processor.process(images=ip_img, text=query_text)
    # Move inputs to the correct device and create a batch of size 1
    inputs = {k: v.to(model.device).unsqueeze(0) for k, v in inputs.items()}

    # Generate output; maximum 200 new tokens; stop generation when <|endoftext|> is generated
    output = model.generate_from_batch(
    inputs, GenerationConfig(max_new_tokens=500, stop_strings="<|endoftext|>"), tokenizer=processor.tokenizer
    )

    # Only get generated tokens; decode them to text
    generated_tokens = output[0, inputs["input_ids"].size(1) :]
    generated_text = processor.tokenizer.decode(generated_tokens, skip_special_tokens=True)

    # Print the generated text
    print(generated_text)


    # format_chat = f"Previous Query: {query_text}\nYour Previous Response: {generated_text}"
    format_chat = {
        "Previous Query": input_query,
        "Previous Model Response": generated_text
    }

    print("----------------------------------------------")
    print(format_chat)
    print("------------------------------------------------")

    chat_history = save_history(chat_history, save_text = format_chat, remove_prev_str = False)



def save_history(chat_history, save_text: str = '', remove_prev_str: bool = False, remove_len: int = 1):
    
    chat_history = [save_text]+chat_history
    
    if remove_prev_str:
        for _ in range(remove_len):
            chat_history.pop(0)
    
    return chat_history