various-fixes (#4)

- Updated readme (62753ff252dc290433fb40851374cfe46efb17f6)
- Script to run model updated (5cdaab30461cc485f414053fd4839ab63399b5e7)

README.md

@@ -10,34 +10,21 @@ ICELANDIC GPT-SW3 FOR SPELL AND GRAMMAR CHECKING
 
 This is a model for correcting spelling and grammar errors in Icelandic text. It is a GPT-SW3 model (https://huggingface.co/AI-Sweden-Models/gpt-sw3-6.7b) finetuned on Icelandic and particularly on the spell and grammar checking task.
 
-Provided here is the model along with a script for running it locally.
+Provided here is the model along with a script for running it through a Hugging Face endpoint. An authorized Hugging Face API key is required to do so. Once you have retrieved an API key and it has been authorized, add it to you environment as "HF_API_KEY".
 
-To run the script you will need a python3 environment. Install the required dependencies by running
+To run the model you will need a python3 environment. Install the required dependencies by running
 
 > pip install -r requirements.txt
 
 The current version of transformers includes a bug which has to be fixed in the user's environment before the model can be run. To fix it, change "gpt-sw3-7b" in line no. 138 in transformers/models/gpt_sw3/tokenization_gpt_sw3.py to "gpt-sw3-6.7b".
 
-After that you can run the script with an input file consisting of text to correct. 
-
-The model is fine-tuned on the following three tasks:
+The model is fine-tuned on the following three tasks. Output examples for each task are shown in ./example_outputs.
   - Task 1: The model evaluates one text with regards to e.g. grammar and spelling, and returns all errors in the input text as a list, with their position in the text and their corrections.
   - Task 2: The model evaluates two texts and chooses which one is better with regards to e.g. grammar and spelling.
   - Task 3: The model evaluates one text with regards to e.g. grammar and spelling, and returns a corrected version of the text.
 
-The script which runs the model takes the following three arguments:
-  - --task: A number (1-3) representing the intended task. The script includes prompts for each task.
-  - --input-file: A file containing text to be evaluated. The format of the input file differs between tasks, and is described further below.
-  - --output-file: A path to a desired output file to be created by the script. The format of the file differs between tasks, and is described further below.
-
-An input file for tasks 1 and 3 should be a .txt file consisting of texts per line. An example of both files can be found under ./example_inputs. 
-An input file for task 2 should be a .jsonl file, where each line is a dictionary object showing two texts. Keys in the dictionary are "a" and "b" and texts to be evaluated are their values. An example of this file can be found under ./example_inputs.
-
-All output files are .txt files and output examples for each task are shown in ./example_outputs. An output file for task 1 shows each text which was evaluated, followed by a list of corrections. Text outputs are separated by an empty line. An output file for task 2 shows 'A' or 'B' for which text is preferred, one choice per line. An output file for task 3 shows the corrected text, one text per line.
-
-Run the script with
-
-> python run_model.py --task 3 --input-file example_inputs/task3_example.txt --output-file example_outputs/task3_example.txt
+Run the model with
 
-The script we provide runs in CPU-only mode and should work on most systems that have enough RAM to load the model. Users that wish to accelerate their corrections with specialized hardware (eg GPUs) will need to install appropriate support packages for their hardware. We refer to the PyTorch documentation: https://pytorch.org/get-started/locally/ . After the extra packages are installed, add the `device` parameter to the pipeline constructor. See the HuggingFace documentation (https://huggingface.co/docs/transformers/main_classes/pipelines) for more details.
+> python run_model.py
 
+Input text(s) and the task type need to be specified in the script.

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)