diff options
Diffstat (limited to 'inference2.py')
| -rw-r--r-- | inference2.py | 179 |
1 files changed, 0 insertions, 179 deletions
diff --git a/inference2.py b/inference2.py deleted file mode 100644 index 67e633a..0000000 --- a/inference2.py +++ /dev/null @@ -1,179 +0,0 @@ -#!/usr/bin/env python3 - -import argparse -import datasets -import json -import os -import requests -import sys -import torch -from glob import glob -from PIL import Image -from transformers import AutoProcessor, AutoModelForImageTextToText - - -def clean_str(s): - return s.removeprefix("```json").removesuffix("```").replace("\n", "").strip() - - -def captioning_prompt(image): - return [ - { - "role": "system", - "content": [{"type": "text", "text": "You are a professional English-German translator and also a renowned photography critic."}] - }, - { - "role": "user", - "content": [ - {"type": "image", "image": image}, - {"type": "text", "text": "Write a detailed caption for this image in a single sentence. Translate the caption into German. The output needs to be JSON, the keys being 'English' and 'German' for the respective captions. Only output the JSON, nothing else."} - ] - } - ] - - -def captioning_prompt_with_source(image, source): - prompt = captioning_prompt(image) - prefix = json.dumps({"English": source}).removesuffix("}") + ', "German": "' - prompt.append({"role": "assistant", "content": [{"type": "text", "text": prefix}]}) - - return prompt - - -def translation_prompt(source): - return [ - { - "role": "system", - "content": [{"type": "text", "text": "You are a professional English-German translator."}] - }, - { - "role": "user", - "content": [ - {"type": "text", "text": f"Translate the following caption into German. The output needs to be JSON, the only being 'Translation' for the translation. Only output the JSON, nothing else. Caption: {source}"} - ] - } - ] - - -def make_inputs(processor, - messages, - device): - return processor.apply_chat_template( - messages, - add_generation_prompt=True, - tokenize=True, - return_dict=True, - return_tensors="pt" - ).to(device, dtype=torch.bfloat16) - - -def generate_and_parse(model, - processor, - messages, - args, - example_id=None): - sys.stderr.write(f"Processing {example_id=}\n") - inputs = make_inputs(processor, messages, model.device) - input_len = inputs["input_ids"].shape[-1] - - stop_token_ids = [processor.tokenizer.eos_token_id, processor.tokenizer.convert_tokens_to_ids("<end_of_turn>")] - - with torch.inference_mode(): - generation = model.generate( - **inputs, - max_new_tokens=args.max_new_tokens, - do_sample=not args.do_not_sample, - temperature=args.temperature, - top_p=args.top_p, - top_k=args.top_k, - eos_token_id=stop_token_ids, - disable_compile=True, - ) - - output_tokens = generation[0][input_len:] - output_text = clean_str(processor.decode(output_tokens, skip_special_tokens=True)) - - try: - return json.loads(output_text) - except Exception: - if example_id is not None: - sys.stderr.write( - f"Error loading JSON from string '{output_text}' for id={example_id}\n" - ) - else: - sys.stderr.write( - f"Error loading JSON from string '{output_text}'\n" - ) - return output_text - - -def main(): - parser = argparse.ArgumentParser() - parser.add_argument("--model", default="google/gemma-3-4b-it", type=str) - parser.add_argument("--attention-implementation", default="eager", type=str) - parser.add_argument("--lora-adapter", default=None, type=str) - parser.add_argument("--mode", choices=["from_scratch", "with_prefix", "translate"], type=str, required=True) - parser.add_argument("--dataset", default="asdf2k/caption_translation", type=str) - parser.add_argument("--data-subset", choices=["train", "dev", "test"], default="test", type=str) - parser.add_argument("--max-new-tokens", default=300, type=int) - parser.add_argument("--top-p", default=1.0, type=int) - parser.add_argument("--top-k", default=50, type=int) - parser.add_argument("--temperature", default=0.8, type=int) - parser.add_argument("--do-not-sample", action="store_true") - args = parser.parse_args() - - model = AutoModelForImageTextToText.from_pretrained( - args.model, - device_map="cuda", - dtype=torch.bfloat16, - attn_implementation=args.attention_implementation, - ).eval() - processor = AutoProcessor.from_pretrained(args.model, use_fast=True) - - if args.lora_adapter: - from peft import PeftModel - model = PeftModel.from_pretrained(model, args.lora_adapter) - - dataset = datasets.load_dataset(args.dataset)[args.data_subset] - - if args.mode == "from_scratch": # Generate caption & translation from scratch - for x in dataset: - output = generate_and_parse( - model, - processor, - captioning_prompt(x["image"]), - args, - example_id=x["id"], - ) - print(f"{x['id']}\t{json.dumps(output)}") - - elif args.mode == "translate": # Generate German translation given English source - for x in dataset: - input_data = json.loads(x["assistant"]) - output = generate_and_parse( - model, - processor, - translation_prompt(input_data["English"]), - args, - example_id=x["id"], - ) - output = {"English": input_data["English"], "German": output["Translation"]} - print(f"{x['id']}\t{json.dumps(output)}") - - elif args.mode == "with_prefix": # Generate German translation given English caption and image - for x in dataset: - assistant_output_as_input = json.loads(x["assistant"]) - output = generate_and_parse( - model, - processor, - captioning_prompt_with_source(x["image"], assistant_output_as_input["English"]), - args, - example_id=x["id"], - ) - print(f"{x['id']}\t{json.dumps(output)}") - else: - sys.stderr.write(f"Unkown mode '{args.mode}'") - - -if __name__ == "__main__": - main() |