import torch
from transformers import (
    AutoTokenizer,
    AutoModelForCausalLM,
)

repo_id = "summerMC/action1"

tokenizer = AutoTokenizer.from_pretrained(
    repo_id,
    trust_remote_code=True,
)

inference_model = AutoModelForCausalLM.from_pretrained(
    repo_id,
    trust_remote_code=True,
    device_map="auto",
)

inference_model.eval()


@torch.no_grad()
def choose_action_greedy(
    model,
    tokenizer,
    goal,
    obs,
    history=None,
    max_new_tokens=64,
):
    if history is None:
        history = []

    prompt = (
        f"Goal: {goal}\n"
        f"Observation:\n{obs}\n"
    )

    if history:
        prompt += "History:\n"
        for item in history:
            prompt += f"{item}\n"

    prompt += "Action:"

    input_ids = tokenizer(
        prompt,
        return_tensors="pt",
        add_special_tokens=False,
    )["input_ids"].to(model.device)

    generated = []

    for _ in range(max_new_tokens):
        if input_ids.shape[1] >= model.config.max_seq_len:
            break

        output = model(
            input_ids=input_ids,
        )

        next_token = torch.argmax(
            output.logits[:, -1, :],
            dim=-1,
            keepdim=True,
        )

        token_id = next_token.item()

        if (
            tokenizer.eos_token_id is not None
            and token_id == tokenizer.eos_token_id
        ):
            break

        generated.append(token_id)

        input_ids = torch.cat(
            [input_ids, next_token],
            dim=1,
        )

        text = tokenizer.decode(
            generated,
            skip_special_tokens=True,
        )

        if "\n" in text:
            return text.split("\n", 1)[0].strip()

    return tokenizer.decode(
        generated,
        skip_special_tokens=True,
    ).strip()


test_goal = "Search for OpenAI"

test_obs = '''[textbox_1] textbox "Search"
[button_1] button "Search"'''

action = choose_action_greedy(
    inference_model,
    tokenizer,
    test_goal,
    test_obs,
    history=[],
)

print("Goal:", test_goal)
print("Action:", action)