Script to run model updated

run_model.py

@@ -1,188 +1,34 @@
-"""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,
+"""
+Script for running the model using the Hugging Face endpoint. An authorized Hugging Face API key is required.
+"""
+
+import requests
+import os
+
+API_URL = "https://otaf5w2ge8huxngl.eu-west-1.aws.endpoints.huggingface.cloud"
+# Set your Hugging Face API key as an environment variable
+api_key = os.environ.get("HF_API_KEY")
+headers = {
+    "Accept": "application/json",
+    "Authorization": f"Bearer {api_key}",
+    "Content-Type": "application/json",
 }
-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,
-        choices=(1,2,3),
-        required=True,
-        help="The task type (1, 2, or 3)",
-    )
-    parser.add_argument(
-        "--input-file",
-        type=str,
-        required=True,
-        help="The path to the input file with data to be corrected",
-    )
-    parser.add_argument(
-        "--output-file",
-        type=str,
-        required=True,
-        help="The path to the output file where the corrected data will be saved",
-    )
-    args = parser.parse_args()
+def query(payload):
+    response = requests.post(API_URL, headers=headers, json=payload)
+    return response.json()
 
-    model_path = "."
-    tokenizer_name = "AI-Sweden-Models/gpt-sw3-6.7b"
 
-    set_seed(42)
-    run_model_on_data(model_path, tokenizer_name, args)
+output = query(
+    {
+        "inputs": "",  # Can be left empty.
+        "input_a": "<text A>",  # Required for all tasks.
+        "input_b": "<text B>",  # Required for task 2 but not for task 1 or 3.
+        "task": 1 | 2 | 3,  # Choose the task number.
+        "parameters": {
+            # Can be left empty
+        },
+    }
+)
+print(output)