| | import argparse
|
| | from transformers import AutoTokenizer, AutoModelForCausalLM, StoppingCriteria
|
| | import torch
|
| | import os
|
| | import json
|
| | from tqdm import tqdm
|
| | import shortuuid
|
| |
|
| | from llava.conversation import default_conversation
|
| | from llava.utils import disable_torch_init
|
| |
|
| |
|
| |
|
| | class KeywordsStoppingCriteria(StoppingCriteria):
|
| | def __init__(self, keywords, tokenizer, input_ids):
|
| | self.keywords = keywords
|
| | self.tokenizer = tokenizer
|
| | self.start_len = None
|
| | self.input_ids = input_ids
|
| |
|
| | def __call__(self, output_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:
|
| | if self.start_len is None:
|
| | self.start_len = self.input_ids.shape[1]
|
| | else:
|
| | outputs = self.tokenizer.batch_decode(output_ids[:, self.start_len:], skip_special_tokens=True)[0]
|
| | for keyword in self.keywords:
|
| | if keyword in outputs:
|
| | return True
|
| | return False
|
| |
|
| |
|
| | @torch.inference_mode()
|
| | def eval_model(model_name, questions_file, answers_file):
|
| |
|
| | disable_torch_init()
|
| | model_name = os.path.expanduser(model_name)
|
| | tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=False)
|
| | model = AutoModelForCausalLM.from_pretrained(model_name,
|
| | torch_dtype=torch.float16).cuda()
|
| |
|
| |
|
| | ques_file = open(os.path.expanduser(questions_file), "r")
|
| | ans_file = open(os.path.expanduser(answers_file), "w")
|
| | for i, line in enumerate(tqdm(ques_file)):
|
| | idx = json.loads(line)["question_id"]
|
| | qs = json.loads(line)["text"]
|
| | cat = json.loads(line)["category"]
|
| | conv = default_conversation.copy()
|
| | conv.append_message(conv.roles[0], qs)
|
| | prompt = conv.get_prompt()
|
| | inputs = tokenizer([prompt])
|
| | input_ids = torch.as_tensor(inputs.input_ids).cuda()
|
| | stopping_criteria = KeywordsStoppingCriteria([conv.sep], tokenizer, input_ids)
|
| | output_ids = model.generate(
|
| | input_ids,
|
| | do_sample=True,
|
| | use_cache=True,
|
| | temperature=0.7,
|
| | max_new_tokens=1024,
|
| | stopping_criteria=[stopping_criteria])
|
| | outputs = tokenizer.batch_decode(output_ids, skip_special_tokens=True)[0]
|
| | try:
|
| | index = outputs.index(conv.sep, len(prompt))
|
| | except ValueError:
|
| | outputs += conv.sep
|
| | index = outputs.index(conv.sep, len(prompt))
|
| |
|
| | outputs = outputs[len(prompt) + len(conv.roles[1]) + 2:index].strip()
|
| | ans_id = shortuuid.uuid()
|
| | ans_file.write(json.dumps({"question_id": idx,
|
| | "text": outputs,
|
| | "answer_id": ans_id,
|
| | "model_id": model_name,
|
| | "metadata": {}}) + "\n")
|
| | ans_file.flush()
|
| | ans_file.close()
|
| |
|
| | if __name__ == "__main__":
|
| | parser = argparse.ArgumentParser()
|
| | parser.add_argument("--model-name", type=str, default="facebook/opt-350m")
|
| | parser.add_argument("--question-file", type=str, default="tables/question.jsonl")
|
| | parser.add_argument("--answers-file", type=str, default="answer.jsonl")
|
| | args = parser.parse_args()
|
| |
|
| | eval_model(args.model_name, args.question_file, args.answers_file)
|
| |
|
