"""This script runs the trained model on data and saves the predictions to a file.""" from transformers import AutoModelForCausalLM, AutoTokenizer import torch import logging import random import tqdm import json import argparse # Set the logging level to info logging.basicConfig(level=logging.INFO) # Set the device to GPU if available device = torch.device("cuda" if torch.cuda.is_available() else "cpu") logging.info(f"Device: {device}") # Prompts for the different tasks START_PROMPT_TASK1 = "Hér er texti sem ég vil að þú skoðir vel og vandlega. Þú skalt skoða hvert einasta orð, orðasamband, og setningu og meta hvort þér finnist eitthvað athugavert, til dæmis hvað varðar málfræði, stafsetningu, skringilega merkingu og svo framvegis.\nHér er textinn:\n\n" END_PROMPT_TASK1 = "Sérðu eitthvað sem mætti betur fara í textanum? Búðu til lista af öllum slíkum tilvikum þar sem hver lína tilgreinir hver villan er, hvar hún er, og hvað væri gert í staðinn fyrir villuna.\n\n" START_PROMPT_TASK2 = "Hér er texti sem ég vil að þú skoðir vel og vandlega. Þú skalt skoða hvert einasta orð, orðasamband, og setningu og meta hvort þér finnist eitthvað athugavert, til dæmis hvað varðar málfræði, stafsetningu, skringilega merkingu og svo framvegis.Ég er með tvær útgáfur af textanum, A og B, og önnur þeirra gæti verið betri en hin á einhvern hátt, t.d. hvað varðar stafsetningu, málfræði o.s.frv.\nHér er texti A:\n\n" MIDDLE_PROMPT_TASK2 = "Hér er texti B:\n\n" END_PROMPT_TASK2 = "Hvorn textann líst þér betur á?\n\n" START_PROMPT_TASK3 = "Hér er texti sem ég vil að þú skoðir vel og vandlega. Þú skalt skoða hvert einasta orð, orðasamband, og setningu og meta hvort þér finnist eitthvað athugavert, til dæmis hvað varðar málfræði, stafsetningu, skringilega merkingu og svo framvegis.\nHér er textinn:\n\n" END_PROMPT_TASK3 = "Reyndu nú að laga textann þannig að hann líti betur út, eins og þér finnst best við hæfi.\n\n" START_PROMPT_TASK = { 1: START_PROMPT_TASK1, 2: START_PROMPT_TASK2, 3: START_PROMPT_TASK3, } END_PROMPT_TASK = {1: END_PROMPT_TASK1, 2: END_PROMPT_TASK2, 3: END_PROMPT_TASK3} SEP = "\n\n" def set_seed(seed): """Set the random seed for reproducibility.""" torch.manual_seed(seed) if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed) torch.backends.cudnn.deterministic = True torch.backends.cudnn.benchmark = False random.seed(seed) def tokenize_data(tokenizer, data, task, max_length): """Tokenize the data and return the input_ids and attention_mask.""" tokenized_start = tokenizer(START_PROMPT_TASK[task])["input_ids"] tokenized_end = tokenizer(END_PROMPT_TASK[task])["input_ids"] if task == 2: tokenized_middle = tokenizer(MIDDLE_PROMPT_TASK2)["input_ids"] # Tokenize the data tokenized_data = [] if task == 1 or task == 3: for sentence in data: tokenized_sentence = tokenizer(sentence + SEP)["input_ids"] # Concatenate the tokenized data concatted_data = ( [tokenizer.bos_token_id] + tokenized_start + tokenized_sentence + tokenized_end ) # Truncate the data concatted_data = concatted_data[:max_length] tokenized_data.append(concatted_data) elif task == 2: for line in data: data_a = line["a"] data_b = line["b"] tokenized_sentence_a = tokenizer(data_a + SEP)["input_ids"] tokenized_sentence_b = tokenizer(data_b + SEP)["input_ids"] # Concatenate the tokenized data concatted_data = ( [tokenizer.bos_token_id] + tokenized_start + tokenized_sentence_a + tokenized_middle + tokenized_sentence_b + tokenized_end ) # Truncate the data concatted_data = concatted_data[:max_length] tokenized_data.append(concatted_data) return tokenized_data def run_model_on_data(model_path, tokenizer_name, arguments): """Run the model on the data and save the predictions to a file.""" # Load the model and tokenizer model = AutoModelForCausalLM.from_pretrained(model_path, torch_dtype=torch.bfloat16) model.to(device) model.eval() tokenizer = AutoTokenizer.from_pretrained(tokenizer_name) # Load the data if arguments.task == 1 or arguments.task == 3: with open(arguments.input_file, "r") as file: data = file.read().splitlines() elif arguments.task == 2: with open(arguments.input_file, "r") as file: data = file.read().splitlines() data = [json.loads(line) for line in data] # Tokenize the data data_tokenized = tokenize_data( tokenizer, data, arguments.task, tokenizer.model_max_length ) logging.info(f"Number of examples: {len(data_tokenized)}") # Run the model on the data predictions = [] progress_bar = tqdm.tqdm(total=len(data_tokenized), desc="Running model on data") for input_ids in data_tokenized: progress_bar.update(1) # Generate the predictions with torch.cuda.amp.autocast(dtype=torch.bfloat16): input_ids_tensor = torch.tensor(input_ids).unsqueeze(0).to(device) output = model.generate( input_ids=input_ids_tensor, max_new_tokens=500, num_return_sequences=1 ) # Only get the part of the prediction that was generated prediction = tokenizer.decode( output[0][len(input_ids) :], skip_special_tokens=True ) predictions.append(prediction) progress_bar.close() # Save the predictions to a file with open(arguments.output_file, "w") as file: if arguments.task == 1: # We want to include the original text in the output file counter = 0 for prediction in predictions: file.write(data[counter] + "\n") file.write(prediction.split("\n\n")[0] + "\n\n") counter += 1 else: for prediction in predictions: file.write(prediction.split("\n\n")[0] + "\n") logging.info(f"Predictions written to file: {arguments.output_file}") if __name__ == "__main__": # Parse the arguments parser = argparse.ArgumentParser() parser.add_argument("--task", type=int, help="The task type (1, 2, or 3)") parser.add_argument( "--input-file", type=str, help="The path to the input file with data to be corrected", ) parser.add_argument( "--output-file", type=str, help="The path to the output file where the corrected data will be saved", ) args = parser.parse_args() model_path = "./gpt-sw3-model" tokenizer_name = "AI-Sweden-Models/gpt-sw3-6.7b" set_seed(42) run_model_on_data(model_path, tokenizer_name, args)