File size: 2,086 Bytes
0aabfc1 b5373b0 c49b61b 0aabfc1 c49b61b b5373b0 c49b61b b5373b0 c49b61b b5373b0 c49b61b b5373b0 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 | 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) |