format and clean code

app.py

@@ -1,9 +1,9 @@
-import gradio as gr
 import os
+
+import gradio as gr
 from openai import OpenAI
-from generate_prompt import construct_generic_prompt, recommend_config
 
-# Define available tasks and their corresponding datasets
+from generate_prompt import construct_generic_prompt, recommend_config
 
 QA = "QA"
 SUMMARIZATION = "Summarization"
@@ -14,21 +14,59 @@ tasks_datasets = {
     QA: ["XQuad", "Indicqa"],
     SUMMARIZATION: ["XLSum", "HeSum"],
     NLI: ["XNLI"],
-    NER: ["MasakaNER", "WikiANN"]
+    NER: ["MasakaNER", "WikiANN"],
 }
 
 # List of all languages
 languages = [
-    "English", "Spanish", "French", "German", "Chinese", "Japanese", "Korean", "Italian",
-    "Portuguese", "Russian", "Arabic", "Hindi", "Bengali", "Turkish", "Vietnamese", "Polish",
-    "Dutch", "Indonesian", "Malay", "Thai", "Greek", "Swedish", "Hungarian", "Finnish",
-    "Danish", "Norwegian", "Hebrew", "Czech", "Slovak", "Bulgarian", "Romanian", "Serbian",
-    "Croatian", "Ukrainian", "Lithuanian", "Latvian", "Estonian", "Filipino", "Icelandic",
-    "Irish", "Welsh", "Maltese", "Swahili", "Zulu", "Afrikaans"
+    "English",
+    "Spanish",
+    "French",
+    "German",
+    "Chinese",
+    "Japanese",
+    "Korean",
+    "Italian",
+    "Portuguese",
+    "Russian",
+    "Arabic",
+    "Hindi",
+    "Bengali",
+    "Turkish",
+    "Vietnamese",
+    "Polish",
+    "Dutch",
+    "Indonesian",
+    "Malay",
+    "Thai",
+    "Greek",
+    "Swedish",
+    "Hungarian",
+    "Finnish",
+    "Danish",
+    "Norwegian",
+    "Hebrew",
+    "Czech",
+    "Slovak",
+    "Bulgarian",
+    "Romanian",
+    "Serbian",
+    "Croatian",
+    "Ukrainian",
+    "Lithuanian",
+    "Latvian",
+    "Estonian",
+    "Filipino",
+    "Icelandic",
+    "Irish",
+    "Welsh",
+    "Maltese",
+    "Swahili",
+    "Zulu",
+    "Afrikaans",
 ]
 
 
-
 def get_datasets(task):
     return tasks_datasets.get(task, [])
 
@@ -39,16 +77,25 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
 
     with gr.Accordion(label="Task Details", open=True):
         with gr.Row():
-            task = gr.Dropdown(label="Task", choices=list(tasks_datasets.keys()), value=QA)
-            language = gr.Dropdown(label="Source Language", choices=languages, value="English")
-            model_type = gr.Dropdown(label="Model Type", choices=["Multilingual", "Standard"], value='Multilingual')
+            task = gr.Dropdown(
+                label="Task", choices=list(tasks_datasets.keys()), value=QA
+            )
+            language = gr.Dropdown(
+                label="Source Language", choices=languages, value="English"
+            )
+            model_type = gr.Dropdown(
+                label="Model Type",
+                choices=["Multilingual", "Standard"],
+                value="Multilingual",
+            )
         config_recommendation = gr.Button("Recommend Configuration")
         with gr.Row():
-            config_prompt = gr.Textbox(label="Recommended Configuration", interactive=False,
-                                placeholder="Recommended Configuration for this scenerio")
+            config_prompt = gr.Textbox(
+                label="Recommended Configuration",
+                interactive=False,
+                placeholder="Recommended Configuration for this scenerio",
+            )
     with gr.Row():
-        # examples_selection = gr.Dropdown(["English", "Source"], label="examples", value='English')
-        # output_selection = gr.Dropdown(["English", "Source"], label="output", value='English')
         with gr.Accordion(label="Prompt Template", open=True):
             with gr.Column(scale=2):
                 # Set the same background style across all components
@@ -56,16 +103,41 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
                     instruction = gr.Textbox(label="Instruction")
                     with gr.Row(variant="panel"):
                         zero_shot = gr.Checkbox(label="Zero Shot Setting", value=False)
-                        with gr.Accordion("Few Shot - Select Type of Examples ", open=False, visible=True) as few_shot:
-                            dataset = gr.Dropdown(label="Dataset", choices=tasks_datasets[QA], value="XlSum")
-                            num_examples = gr.Slider(label="Number of examples in context", minimum=1, maximum=10, step=1,
-                                                value=3)
+                        with gr.Accordion(
+                            "Few Shot - Select Type of Examples ",
+                            open=False,
+                            visible=True,
+                        ) as few_shot:
+                            dataset = gr.Dropdown(
+                                label="Dataset",
+                                choices=tasks_datasets[QA],
+                                value="XlSum",
+                            )
+                            num_examples = gr.Slider(
+                                label="Number of examples in context",
+                                minimum=1,
+                                maximum=10,
+                                step=1,
+                                value=3,
+                            )
                     with gr.Row(equal_height=True, variant="panel"):
-                        with gr.Accordion(label="Language Component Selection", open=False):
-                            prefix_selection = gr.Dropdown(["English", "Source"], label="instruction", value='English')
-                            context_selection = gr.Dropdown(["English", "Source"], label="context", value='English')
-                            examples_selection = gr.Dropdown(["English", "Source"], label="examples", value='English')
-                            output_selection = gr.Dropdown(["English", "Source"], label="output", value='English')
+                        with gr.Accordion(
+                            label="Language Component Selection", open=False
+                        ):
+                            prefix_selection = gr.Dropdown(
+                                ["English", "Source"],
+                                label="instruction",
+                                value="English",
+                            )
+                            context_selection = gr.Dropdown(
+                                ["English", "Source"], label="context", value="English"
+                            )
+                            examples_selection = gr.Dropdown(
+                                ["English", "Source"], label="examples", value="English"
+                            )
+                            output_selection = gr.Dropdown(
+                                ["English", "Source"], label="output", value="English"
+                            )
                             # Accordion for Few Shot example selection
                         with gr.Accordion(label="Prompt Input Data", open=False):
                             question = gr.Textbox(label="Question", visible=True)
@@ -78,87 +150,145 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
     generate_button = gr.Button("Generate Prompt")
 
     with gr.Row():
-        prompt = gr.Textbox(label="Generated Prompt", interactive=False, placeholder="Generated prompt will appear here.")
-
+        prompt = gr.Textbox(
+            label="Generated Prompt",
+            interactive=False,
+            placeholder="Generated prompt will appear here.",
+        )
 
     def update_datasets(selected_task):
         return gr.Dropdown(choices=get_datasets(selected_task))
 
-
     def toggle_task_inputs(selected_task):
         if selected_task == QA:
             return (
-                gr.update(visible=True), gr.update(visible=True), gr.update(visible=False),
-                gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
+                gr.update(visible=True),
+                gr.update(visible=True),
+                gr.update(visible=False),
+                gr.update(visible=False),
+                gr.update(visible=False),
+                gr.update(visible=False),
             )
         elif selected_task == SUMMARIZATION:
             return (
-                gr.update(visible=False), gr.update(visible=False), gr.update(visible=True),
-                gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
+                gr.update(visible=False),
+                gr.update(visible=False),
+                gr.update(visible=True),
+                gr.update(visible=False),
+                gr.update(visible=False),
+                gr.update(visible=False),
             )
         elif selected_task == NER:
             return (
-                gr.update(visible=False), gr.update(visible=False), gr.update(visible=False),
-                gr.update(visible=True), gr.update(visible=False), gr.update(visible=False)
+                gr.update(visible=False),
+                gr.update(visible=False),
+                gr.update(visible=False),
+                gr.update(visible=True),
+                gr.update(visible=False),
+                gr.update(visible=False),
             )
         else:
             return (
-                gr.update(visible=False), gr.update(visible=False), gr.update(visible=False),
-                gr.update(visible=False), gr.update(visible=True), gr.update(visible=True)
+                gr.update(visible=False),
+                gr.update(visible=False),
+                gr.update(visible=False),
+                gr.update(visible=False),
+                gr.update(visible=True),
+                gr.update(visible=True),
             )
 
-
     def toggle_num_examples(zero_shot_value):
         # If zero_shot is True, hide the num_examples slider
         return gr.update(visible=not zero_shot_value)
 
     def update_language_selection(language):
-        return gr.update(choices=list({'English', language})), gr.update(choices=list({'English', language})), gr.update(choices=list({'English', language})), gr.update(choices=list({'English', language}))
+        return (
+            gr.update(choices=list({"English", language})),
+            gr.update(choices=list({"English", language})),
+            gr.update(choices=list({"English", language})),
+            gr.update(choices=list({"English", language})),
+        )
 
-    def generatePrompt(instruction, num_examples, zero_shot,
-            task, selected_language, dataset, prefix_selection, context_selection, examples_selection, output_selection,
-            text,  question, context, sentence, hypothesis, premise):
+    def generatePrompt(
+        instruction,
+        num_examples,
+        zero_shot,
+        task,
+        selected_language,
+        dataset,
+        prefix_selection,
+        context_selection,
+        examples_selection,
+        output_selection,
+        text,
+        question,
+        context,
+        sentence,
+        hypothesis,
+        premise,
+    ):
 
-        config = {'prefix': str.lower(prefix_selection), 'input': str.lower(context_selection), 'context': str.lower(examples_selection), 'output': str.lower(output_selection)}
+        config = {
+            "prefix": str.lower(prefix_selection),
+            "input": str.lower(context_selection),
+            "context": str.lower(examples_selection),
+            "output": str.lower(output_selection),
+        }
 
         if task == QA:
             text_example = {
-                'context': context,
-                'question': question,
+                "context": context,
+                "question": question,
             }
         elif task == SUMMARIZATION:
             text_example = {
-                'text': text,
+                "text": text,
             }
         elif task == NER:
-            text_example = {
-                'tokens': sentence,
-                'ner_tags': ''
-            }
+            text_example = {"tokens": sentence, "ner_tags": ""}
         else:
-            text_example = {
-                'hypothesis': hypothesis,
-                'premise': premise
-            }
+            text_example = {"hypothesis": hypothesis, "premise": premise}
 
-        prompt = construct_generic_prompt(task, instruction, text_example, zero_shot, num_examples, selected_language, dataset, config)
+        prompt = construct_generic_prompt(
+            task,
+            instruction,
+            text_example,
+            zero_shot,
+            num_examples,
+            selected_language,
+            dataset,
+            config,
+        )
 
         return prompt
 
-
-    def respond(message, openai_key, url, chat_history, model, config_input, config_prefix, config_context,
-                config_output, task, dataset, language, num_examples, zero_shot):
+    def respond(
+        message,
+        openai_key,
+        url,
+        chat_history,
+        model,
+        config_input,
+        config_prefix,
+        config_context,
+        config_output,
+        task,
+        dataset,
+        language,
+        num_examples,
+        zero_shot,
+    ):
         os.environ["OPENAI_API_KEY"] = openai_key
         client = OpenAI()
 
         config = {
             "input": config_input,
             "prefix": config_prefix,
-            "context": config_context.split(', '),
+            "context": config_context.split(", "),
             "output": config_output,
             "language": language,
             "num_examples": num_examples,
-            "zero_shot": zero_shot
+            "zero_shot": zero_shot,
         }
 
         response = client.chat.completions.create(
@@ -171,7 +301,7 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
                         {"type": "image_url", "image_url": url},
                         {"type": "config", "config": config},
                         {"type": "task", "text": task},
-                        {"type": "dataset", "text": dataset}
+                        {"type": "dataset", "text": dataset},
                     ],
                 },
             ],
@@ -183,37 +313,60 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
         chat_history.append((message, out))
         return "", chat_history
 
-
     # Bind functions to dropdown changes and button click
     # task.change(fn=update_datasets, outputs=dataset)
-    language.change(fn=update_language_selection, inputs=language, outputs=[prefix_selection, context_selection, examples_selection, output_selection])
+    language.change(
+        fn=update_language_selection,
+        inputs=language,
+        outputs=[
+            prefix_selection,
+            context_selection,
+            examples_selection,
+            output_selection,
+        ],
+    )
 
     zero_shot.change(fn=toggle_num_examples, inputs=zero_shot, outputs=few_shot)
     zero_shot.change(fn=toggle_num_examples, inputs=zero_shot, outputs=num_examples)
     task.change(fn=update_datasets, inputs=task, outputs=dataset)
-    task.change(fn=toggle_task_inputs, inputs=task, outputs=[
-        question, context, text, sentence, hypothesis, premise,
-    ])
+    task.change(
+        fn=toggle_task_inputs,
+        inputs=task,
+        outputs=[
+            question,
+            context,
+            text,
+            sentence,
+            hypothesis,
+            premise,
+        ],
+    )
     generate_button.click(
         generatePrompt,
         inputs=[
-            instruction, num_examples, zero_shot,
-            task, language, dataset, prefix_selection, context_selection, examples_selection, output_selection,
-            text, question, context, sentence, hypothesis, premise
-
+            instruction,
+            num_examples,
+            zero_shot,
+            task,
+            language,
+            dataset,
+            prefix_selection,
+            context_selection,
+            examples_selection,
+            output_selection,
+            text,
+            question,
+            context,
+            sentence,
+            hypothesis,
+            premise,
         ],
-        outputs=[prompt]
+        outputs=[prompt],
     )
 
     config_recommendation.click(
-        recommend_config,
-        inputs=[
-            task,
-            language,
-            model_type
-        ],
-        outputs=[config_prompt]
+        recommend_config, inputs=[task, language, model_type], outputs=[config_prompt]
     )
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     demo.launch(share=True)

generate_prompt.py

@@ -1,31 +1,10 @@
-import collections
-import csv
 import enum
-import json
-import logging
-import os
-import re
-import string
-import sys
-import unicodedata
-from typing import Any, Dict, List, NewType, Union
 
-import numpy as np
-import openai
 import pandas as pd
-import requests
-import yaml
-from datasets import Dataset, load_dataset
-from easygoogletranslate import EasyGoogleTranslate
-from langchain.prompts import FewShotPromptTemplate, PromptTemplate
-from tqdm import tqdm
-from yaml.loader import SafeLoader
 
-from tasks import ner, summarization, qa, nli
+from tasks import ner, nli, qa, summarization
 
 
-# from models.model_completion import gpt3x_completion, gemini_completion
-
 class LanguageType(enum.Enum):
     Low = "Low"
     High = "High"
@@ -36,504 +15,6 @@ class ModelType(enum.Enum):
     Multilingual = "Multilingual"
 
 
-def get_entities_gpt3_long(prompt):
-    response = openai.ChatCompletion.create(
-        engine="chatgpt", temperature=0, messages=[{"role": "user", "content": prompt}]
-    )
-    return response["choices"][0]["message"]["content"]
-
-
-def gpt3x_completion(
-        prompt: Union[str, List[Dict[str, str]]],
-) -> str:
-    import os
-    import openai
-    os.environ["OPENAI_API_KEY"] = '07d805ec4fbd484ebc923a3a41e1773d'
-    OPENAI_API_KEY = '07d805ec4fbd484ebc923a3a41e1773d'
-    openai.api_type = "azure"
-    openai.api_base = 'https://hebsum-itaim-uks.openai.azure.com/'
-    openai.api_version = "2023-03-15-preview"
-    openai.api_key = '07d805ec4fbd484ebc923a3a41e1773d'
-
-    def get_entities_chatGPT(final_prompt):
-        response = openai.ChatCompletion.create(
-            engine="gpt35-16k",
-            temperature=0,
-            messages=[
-                {"role": "user", "content": final_prompt}
-            ]
-        )
-        return response['choices'][0]['message']['content']
-
-    return get_entities_chatGPT(final_prompt=prompt)
-
-
-def mixtral_completion(prompt):
-    url = "https://api.together.xyz/v1/chat/completions"
-
-    # Define your Together API key
-    together_api_key = "851cfc39f3d7a246a2342259f5f6fbba4721c6002123365fba2254c9c9c424ad"  # Replace with your actual API key
-
-    # Define the request payload
-    payload = {
-        "temperature": 0,
-        "max_tokens": 30,
-        "model": "mistralai/Mixtral-8x7B-Instruct-v0.1",
-        "messages": [{"role": "user", "content": f"{prompt}"}],
-    }
-
-    # Define request headers
-    headers = {
-        "Authorization": f"Bearer {together_api_key}",
-        "Content-Type": "application/json",
-    }
-
-    # Send POST request
-    response = requests.post(url, json=payload, headers=headers)
-
-    # Check response status
-    if response.status_code == 200:
-        # Print the response content (API output)
-        return response.json()["choices"][0]["message"]["content"]
-    else:
-        # Print error message if request fails
-        print(f"Error: {response.status_code} - {response.text}")
-
-
-XQUAD_LANG2CODES = {
-    "bengali": "bn",
-    "korean": "ko",
-    "swahili": "sw",
-    "english": "en",
-    "indonesian": "id",
-    "arabic": "ar",
-    "finnish": "fi",
-    "telugu": "te",
-    "russian": "ru",
-    "german": "de",
-    "greek": "el",
-    "hindi": "hi",
-    "vietnamese": "vi",
-    "romanian": "ro",
-}
-
-INDICQA_LANG2CODES = {
-    "indicqa": "as",
-    "bengali": "bn",
-    "gujarati": "gu",
-    "hindi": "hi",
-    "kannada": "kn",
-    "malayalam": "ml",
-    "marathi": "mr",
-    "odia": "or",
-    "punjabi": "pa",
-    "tamil": "ta",
-    "telugu": "te",
-    "assamese": "as",
-}
-
-PUNCT = {
-    chr(i)
-    for i in range(sys.maxunicode)
-    if unicodedata.category(chr(i)).startswith("P")
-}.union(string.punctuation)
-WHITESPACE_LANGS = ["en", "es", "hi", "vi", "de", "ar"]
-MIXED_SEGMENTATION_LANGS = ["zh"]
-
-TYDIQA_LANG2CODES = {
-    "bengali": "bn",
-    "korean": "ko",
-    "swahili": "sw",
-    "english": "en",
-    "indonesian": "id",
-    "arabic": "ar",
-    "finnish": "fi",
-    "telugu": "te",
-    "russian": "ru",
-    "assamese": "as",
-    "persian": "fa",
-}
-
-logger = logging.Logger("Xlsum_task")
-LANGUAGE_TO_SUFFIX = {
-    "chinese_simplified": "zh-CN",
-    "french": "fr",
-    "portuguese": "pt",
-    "english": "en",
-    "arabic": "ar",
-    "hindi": "hi",
-    "indonesian": "id",
-    "amharic": "am",
-    "bengali": "bn",
-    "telugu": "te",
-    "burmese": "my",
-    "german": "de",
-    "greek": "el",
-    "tamil": "ta",
-    "assamese": "as",
-    "hindi": "hi",
-    "vietnamese": "vi",
-    "russian": "ru",
-    "telugu": "te",
-    "romanian": "ro",
-    "malayalam": "ml",
-    "persian": "fa",
-}
-
-PARAMS = NewType("PARAMS", Dict[str, Any])
-
-
-def read_parameters(args_path) -> PARAMS:
-    with open(args_path) as f:
-        args = yaml.load(f, Loader=SafeLoader)
-    return args
-
-
-def load_qa_dataset(dataset_name, lang, split, translate_test=False, limit=5):
-    if dataset_name == "indicqa":
-        if split != "train":
-            dataset = load_dataset(
-                "ai4bharat/IndicQA", f"indicqa.{INDICQA_LANG2CODES[lang]}"
-            )[split]
-        else:
-            dataset = load_dataset("squad_v2")[split]
-    elif dataset_name == "xquad":
-        if split != "train":
-            dataset = load_dataset("xquad", f"xquad.{XQUAD_LANG2CODES[lang]}")[
-                "validation"
-            ]
-        else:
-            dataset = load_dataset("squad")[split]
-    elif dataset_name == "tydiqa":
-        dataset = load_dataset("tydiqa", "secondary_task")[split]
-        dataset = dataset.map(
-            lambda example: {"lang": TYDIQA_LANG2CODES[example["id"].split("-")[0]]}
-        )
-        dataset = dataset.filter(lambda example: example["lang"] == lang)
-    elif dataset_name == "mlqa":
-        if split == "train":
-            print("No Training Data for MLQA, switching to validation!")
-            split = "validation"
-        if translate_test:
-            dataset_name = f"mlqa-translate-test.{lang}"
-        else:
-            dataset_name = f"mlqa.{lang}.{lang}"
-
-        dataset = load_dataset("mlqa", dataset_name)[split]
-
-    else:
-        raise NotImplementedError()
-    return dataset.select(np.arange(limit))
-
-
-def construct_prompt(
-        instruction: str,
-        test_example: dict,
-        ic_examples: List[dict],
-        zero_shot: bool,
-        lang: str,
-        config: Dict[Any, Any],
-):
-    example_prompt = PromptTemplate(
-        input_variables=["context", "question", "answers"],
-        template="Context: {context}\nQuestion: {question}\n" "Answers: {answers}",
-    )
-
-    zero_shot_template = (
-            f"""{instruction}""" + "\n<Context>: {context} \n<Question>: {question} " ""
-    )
-
-    prompt = (
-        FewShotPromptTemplate(
-            examples=ic_examples,
-            prefix=instruction,
-            example_prompt=example_prompt,
-            suffix="<Context>: {context} \n<Question>: {question} \nAnswers: ?",
-            input_variables=["question", "context"],
-        )
-        if not zero_shot
-        else PromptTemplate(
-            input_variables=["question", "context"], template=zero_shot_template
-        )
-    )
-
-    label = test_example["answers"]
-    if config["input"] != lang:
-        test_example = _translate_example(
-            example=test_example, src_language=lang, target_language=config["input"]
-        )
-
-    return (
-        prompt.format(
-            question=test_example["question"], context=test_example["context"]
-        ),
-        label,
-    )
-
-
-def dump_metrics(
-        lang: str, config: Dict[str, str], f1: float, em: float, metric_logger_path: str
-):
-    # Check if the metric logger file exists
-    file_exists = os.path.exists(metric_logger_path)
-
-    # Open the CSV file in append mode
-    with open(metric_logger_path, "a", newline="") as f:
-        csvwriter = csv.writer(f, delimiter=",")
-
-        # Write header row if the file is newly created
-        if not file_exists:
-            header = ["Language", "Prefix", "Input", "Context", "Output", "F1", "Em"]
-            csvwriter.writerow(header)
-
-        csvwriter.writerow(
-            [
-                lang,
-                config["prefix"],
-                config["input"],
-                config["context"][0],
-                config["output"],
-                f1,
-                em,
-            ]
-        )
-
-
-def dump_predictions(idx, response, label, response_logger_file):
-    obj = {"q_idx": idx, "prediction": response, "label": label}
-    with open(response_logger_file, "a") as f:
-        f.write(json.dumps(obj, ensure_ascii=False) + "\n")
-
-
-def _translate_instruction(basic_instruction: str, target_language: str) -> str:
-    translator = EasyGoogleTranslate(
-        source_language="en",
-        target_language=LANGUAGE_TO_SUFFIX[target_language],
-        timeout=50,
-    )
-    return translator.translate(basic_instruction)
-
-
-def _translate_prediction_to_output_language(
-        prediction: str, prediction_language: str, output_language: str
-) -> str:
-    translator = EasyGoogleTranslate(
-        source_language=LANGUAGE_TO_SUFFIX[prediction_language],
-        target_language=LANGUAGE_TO_SUFFIX[output_language],
-        timeout=10,
-    )
-    return translator.translate(prediction)
-
-
-def create_instruction(lang: str, expected_output: str):
-    basic_instruction = (
-        "Answer to the <Question> below, based only to the given <Context>, Follow these instructions:\n"
-        "1. The answer should include only words from the given context\n"
-        "2. The answer must include up to 5 words\n"
-        "3. The answer Should be the shortest as possible\n"
-        f"4. The answer must be in {expected_output} only!, not another language!!!"
-    )
-    return (
-        basic_instruction
-        if lang == "english"
-        else _translate_instruction(basic_instruction, target_language=lang)
-    )
-
-
-def _translate_example(
-        example: Dict[str, str], src_language: str, target_language: str
-):
-    translator = EasyGoogleTranslate(
-        source_language=LANGUAGE_TO_SUFFIX[str(src_language).lower()],
-        target_language=LANGUAGE_TO_SUFFIX[str(target_language).lower()],
-        timeout=30,
-    )
-
-    return {
-        "question": translator.translate(example["question"]),
-        "context": translator.translate(example["context"][:2000])
-                   + translator.translate(example["context"][2000:4000])
-                   + translator.translate(example["context"][4000:6000]),
-        "answers": translator.translate(example["answers"][0]),
-    }
-    # except Exception as e:
-    #     print(example["text"])
-    #     print(example["summary"])
-    #     print(e)
-
-
-def choose_few_shot_examples(
-        train_dataset: Dataset,
-        few_shot_size: int,
-        context: List[str],
-        selection_criteria: str,
-        lang: str,
-) -> List[Dict[str, Union[str, int]]]:
-    """Selects few-shot examples from training datasets
-
-    Args:
-        train_dataset (Dataset): Training Dataset
-        few_shot_size (int): Number of few-shot examples
-        selection_criteria (few_shot_selection): How to select few-shot examples. Choices: [random, first_k]
-
-    Returns:
-        List[Dict[str, Union[str, int]]]: Selected examples
-    """
-    selected_examples = []
-
-    example_idxs = []
-    if selection_criteria == "first_k":
-        example_idxs = list(range(few_shot_size))
-    elif selection_criteria == "random":
-        example_idxs = (
-            np.random.choice(len(train_dataset), size=few_shot_size, replace=True)
-            .astype(int)
-            .tolist()
-        )
-
-    ic_examples = [
-        {
-            "question": train_dataset[idx]["question"],
-            "context": train_dataset[idx]["context"],
-            "answers": train_dataset[idx]["answers"]["text"],
-        }
-        for idx in example_idxs
-    ]
-
-    for idx, ic_language in enumerate(context):
-        (
-            selected_examples.append(ic_examples[idx])
-            if ic_language == lang
-            else (
-                selected_examples.append(
-                    _translate_example(
-                        example=ic_examples[idx],
-                        src_language=lang,
-                        target_language=ic_language,
-                    )
-                )
-            )
-        )
-
-    return selected_examples
-
-
-def normalize_answer(s):
-    """Lower text and remove punctuation, articles and extra whitespace."""
-
-    def remove_articles(text):
-        return re.sub(r"\b(a|an|the)\b", " ", text)
-
-    def white_space_fix(text):
-        return " ".join(text.split())
-
-    def remove_punc(text):
-        exclude = set(PUNCT)  # set(string.punctuation)
-        return "".join(ch for ch in text if ch not in exclude)
-
-    def lower(text):
-        return text.lower()
-
-    return white_space_fix(remove_articles(remove_punc(lower(s))))
-
-
-def process_test_example(
-        test_data, config_header, idx, test_example, config, zero_shot, lang, params
-):
-    try:
-        # Your existing code for processing each test example
-        instruction = create_instruction(
-            lang=config["prefix"], expected_output=config["output"]
-        )
-        text_example = {
-            "question": test_example["question"],
-            "context": test_example["context"],
-            "answers": test_example["answers"]["text"],
-        }
-
-        ic_examples = []
-        if not zero_shot:
-            ic_examples = choose_few_shot_examples(
-                train_dataset=test_data,
-                few_shot_size=len(config["context"]),
-                context=config["context"],
-                selection_criteria="random",
-                lang=params["selected_language"],
-            )
-
-        prompt, label = construct_prompt(
-            instruction=instruction,
-            test_example=text_example,
-            ic_examples=ic_examples,
-            zero_shot=zero_shot,
-            lang=lang,
-            config=config,
-        )
-
-        pred = gpt3x_completion(prompt=prompt)
-        print(pred)
-
-        logger.info("Saving prediction to persistent volume")
-        os.makedirs(
-            f"{params['response_logger_root']}/{params['model']}/{lang}", exist_ok=True
-        )
-        dump_predictions(
-            idx=idx,
-            response=pred,
-            label=label,
-            response_logger_file=f"{params['response_logger_root']}/{params['model']}/{lang}/{config_header}.csv",
-        )
-    except Exception as e:
-        # Handle exceptions here
-        print(f"Error processing example {idx}: {e}")
-
-
-def run_one_configuration(selected_language, config, zero_shot, dataset_name, limit=10):
-    test_data = load_qa_dataset(
-        dataset_name=dataset_name,
-        lang=selected_language,
-        split="validation" if dataset_name == "xquad" else "test",
-        limit=limit,
-    )
-
-    for idx, test_example in (pbar := tqdm(enumerate(test_data))):
-        try:
-            instruction = create_instruction(
-                lang=config["prefix"], expected_output=config["output"]
-            )
-            text_example = {
-                "question": test_example["question"],
-                "context": test_example["context"],
-                "answers": test_example["answers"]["text"],
-            }
-
-            ic_examples = []
-            if not zero_shot:
-                ic_examples = choose_few_shot_examples(
-                    train_dataset=test_data,
-                    few_shot_size=len(config["context"]),
-                    context=config["context"],
-                    selection_criteria="random",
-                    lang=selected_language,
-                )
-
-            prompt, label = construct_prompt(
-                instruction=instruction,
-                test_example=text_example,
-                ic_examples=ic_examples,
-                zero_shot=zero_shot,
-                lang=selected_language,
-                config=config,
-            )
-
-            pred = gpt3x_completion(prompt=prompt)
-
-            return pred
-
-        except Exception as e:
-            print(f"Found an exception {e}, continue to the next example")
-            continue
-
 
 QA = "QA"
 SUMMARIZATION = "Summarization"
@@ -541,8 +22,16 @@ NLI = "NLI"
 NER = "NER"
 
 
-def construct_generic_prompt(task, instruction, test_example, zero_shot, num_examples, selected_language, dataset,
-                             config):
+def construct_generic_prompt(
+    task,
+    instruction,
+    test_example,
+    zero_shot,
+    num_examples,
+    selected_language,
+    dataset,
+    config,
+):
     print(task)
     if task == SUMMARIZATION:
         prompt = summarization.construct_prompt(
@@ -588,30 +77,36 @@ def construct_generic_prompt(task, instruction, test_example, zero_shot, num_exa
 
 def _get_language_type(language: str):
     df = pd.read_csv("utils/languages_by_word_count.csv")
-    number_of_words = df[df['Language'] == language]['number of words'].iloc[0]
+    number_of_words = df[df["Language"] == language]["number of words"].iloc[0]
     print(number_of_words)
     return LanguageType.Low if number_of_words < 150276400 else LanguageType.High
 
 
 class Config:
-    def __init__(self, prefix="source", context="source", examples="source", output="source"):
+    def __init__(
+        self, prefix="source", context="source", examples="source", output="source"
+    ):
         self.prefix = prefix
         self.context = context
         self.examples = examples
         self.output = output
 
     def set(self, prefix=None, context=None, examples=None, output=None):
-        if prefix: self.prefix = prefix
-        if context: self.context = context
-        if examples: self.examples = examples
-        if output: self.output = output
+        if prefix:
+            self.prefix = prefix
+        if context:
+            self.context = context
+        if examples:
+            self.examples = examples
+        if output:
+            self.output = output
 
     def to_dict(self):
         return {
-            'instruction': self.prefix,
-            'context': self.context,
-            'examples': self.examples,
-            'output': self.output
+            "instruction": self.prefix,
+            "context": self.context,
+            "examples": self.examples,
+            "output": self.output,
         }
 
 
@@ -622,22 +117,22 @@ def recommend_config(task, lang, model_type):
         if model_type == ModelType.English.value:
             config.set(prefix=lang, context=lang, examples=lang, output=lang)
         else:
-            config.set(prefix='English', context=lang, examples=lang, output=lang)
+            config.set(prefix="English", context=lang, examples=lang, output=lang)
     if task == NER:
         if model_type == ModelType.English.value:
             config.set(prefix=lang, context=lang, examples=lang, output=lang)
         elif language_type == LanguageType.High:
-            config.set(prefix='English', context=lang, examples=lang, output=lang)
+            config.set(prefix="English", context=lang, examples=lang, output=lang)
         else:
-            config.set(prefix='English', context=lang, examples=lang, output='English')
+            config.set(prefix="English", context=lang, examples=lang, output="English")
     if task == NLI:
         if model_type == ModelType.English.value:
             config.set(prefix=lang, context=lang, examples=lang, output=lang)
         elif language_type == LanguageType.High:
-            config.set(prefix='English', context=lang, examples='English')
+            config.set(prefix="English", context=lang, examples="English")
         else:
-            config.set(prefix='English', context='English', examples='English')
+            config.set(prefix="English", context="English", examples="English")
     if task == SUMMARIZATION:
-        config.set(context='English')
+        config.set(context="English")
     print(config.to_dict())
     return config.to_dict()

tasks/ner.py

@@ -1,16 +1,12 @@
-from typing import List, Dict, Union
+from typing import Dict, List, Union
 
 import numpy as np
-from datasets import load_dataset, Dataset
+from datasets import Dataset, load_dataset
 from easygoogletranslate import EasyGoogleTranslate
-from langchain.prompts import PromptTemplate, FewShotPromptTemplate
+from langchain.prompts import FewShotPromptTemplate, PromptTemplate
 
 LANGAUGE_TO_PREFIX = {
-
     "chinese_simplified": "zh-CN",
-    "french": "fr",
-    "portuguese": "pt",
-    "english": "en",
     "arabic": "ar",
     "hindi": "hi",
     "indonesian": "id",
@@ -31,7 +27,6 @@ LANGAUGE_TO_PREFIX = {
     "greek": "el",
     "tamil": "ta",
     "assamese": "as",
-    "vietnamese": "vi",
     "russian": "ru",
     "romanian": "ro",
     "malayalam": "ml",
@@ -39,16 +34,13 @@ LANGAUGE_TO_PREFIX = {
     "bulgarian": "bg",
     "thai": "th",
     "urdu": "ur",
-    "italian": "it",
     "polish": "pl",
     "dutch": "nl",
-    "swedish": "sv",
     "danish": "da",
     "norwegian": "no",
     "finnish": "fi",
     "hungarian": "hu",
     "czech": "cs",
-    "slovak": "sk",
     "ukrainian": "uk",
     "bambara": "bam",
     "ewe": "ewe",
@@ -67,10 +59,9 @@ LANGAUGE_TO_PREFIX = {
     "portuguese": "pt",
     "chinese": "zh",
     "english": "en",
-    "french": "fr"
+    "french": "fr",
 }
 
-
 def _translate_instruction(basic_instruction: str, target_language: str) -> str:
     translator = EasyGoogleTranslate(
         source_language="en",
@@ -104,7 +95,7 @@ def load_wikiann_dataset(lang, split, limit):
 
 
 def _translate_example(
-        example: Dict[str, str], src_language: str, target_language: str
+    example: Dict[str, str], src_language: str, target_language: str
 ):
     translator = EasyGoogleTranslate(
         source_language=LANGAUGE_TO_PREFIX[src_language],
@@ -114,16 +105,16 @@ def _translate_example(
 
     return {
         "tokens": translator.translate(str(example["tokens"])),
-        "ner_tags": translator.translate(str(example["ner_tags"]))
+        "ner_tags": translator.translate(str(example["ner_tags"])),
     }
 
 
 def choose_few_shot_examples(
-        train_dataset: Dataset,
-        few_shot_size: int,
-        context: List[str],
-        selection_criteria: str,
-        lang: str,
+    train_dataset: Dataset,
+    few_shot_size: int,
+    context: List[str],
+    selection_criteria: str,
+    lang: str,
 ) -> List[Dict[str, Union[str, int]]]:
     """Selects few-shot examples from training datasets
 
@@ -150,10 +141,7 @@ def choose_few_shot_examples(
     ic_examples = [train_dataset[idx] for idx in example_idxs]
 
     ic_examples = [
-        {
-            "tokens": ' '.join(example['tokens']),
-            "ner_tags": example['spans']
-        }
+        {"tokens": " ".join(example["tokens"]), "ner_tags": example["spans"]}
         for example in ic_examples
     ]
 
@@ -185,7 +173,7 @@ def construct_prompt(
     config: Dict[str, str],
 ):
     if not instruction:
-        instruction = create_instruction(lang, config['prefix'], config['output'])
+        instruction = create_instruction(lang, config["prefix"], config["output"])
 
     example_prompt = PromptTemplate(
         input_variables=["tokens", "ner_tags"],
@@ -197,8 +185,9 @@ def construct_prompt(
     try:
         test_data = load_wikiann_dataset(lang=lang, split="test", limit=500)
     except Exception as e:
-        raise KeyError(f"{lang} is not supported in 'wikiAnn' dataset, choose supported language in few-shot")
-
+        raise KeyError(
+            f"{lang} is not supported in 'wikiAnn' dataset, choose supported language in few-shot"
+        )
 
     ic_examples = []
     if not zero_shot:

tasks/nli.py

@@ -32,9 +32,7 @@ LANGUAGE_TO_SUFFIX = {
     "spanish": "es",
     "chinese": "zh",
     "greek": "el",
-    "german": "de"
-
-
+    "german": "de",
 }
 
 NUMBER_TO_TAG = {0: "entailment", 1: "neutral", 2: "contradiction"}
@@ -42,9 +40,6 @@ NUMBER_TO_TAG = {0: "entailment", 1: "neutral", 2: "contradiction"}
 PARAMS = NewType("PARAMS", Dict[str, Any])
 
 
-
-
-
 def read_parameters(args_path) -> PARAMS:
     with open(args_path) as f:
         args = yaml.load(f, Loader=SafeLoader)
@@ -278,7 +273,7 @@ def create_instruction(lang: str):
     )
 
 
-def run_one_configuration(params: Optional[PARAMS] = None, zero: bool= False):
+def run_one_configuration(params: Optional[PARAMS] = None, zero: bool = False):
     if not params:
         params = read_parameters("../../parameters.yaml")
 
@@ -320,6 +315,7 @@ def run_one_configuration(params: Optional[PARAMS] = None, zero: bool= False):
     pool.close()
     pool.join()
 
+
 def process_test_example(
     test_data, config_header, idx, test_example, config, zero_shot, lang, params
 ):
@@ -348,7 +344,9 @@ def process_test_example(
             zero_shot=zero_shot,
         )
 
-        pred = get_prediction(prompt=prompt, endpoint_id=7327255438662041600, project_id=16514800572)
+        pred = get_prediction(
+            prompt=prompt, endpoint_id=7327255438662041600, project_id=16514800572
+        )
         print(pred)
 
         os.makedirs(
@@ -367,13 +365,13 @@ def process_test_example(
 
 
 def construct_prompt(
-        instruction: str,
-        test_example: dict,
-        zero_shot: bool,
-        num_examples: int,
-        lang: str,
-        config: Dict[str, str],
-        dataset_name: str = 'xnli'
+    instruction: str,
+    test_example: dict,
+    zero_shot: bool,
+    num_examples: int,
+    lang: str,
+    config: Dict[str, str],
+    dataset_name: str = "xnli",
 ):
     if not instruction:
         print(lang)
@@ -385,13 +383,15 @@ def construct_prompt(
     )
 
     zero_shot_template = (
-            f"""{instruction}""" + "\n Hypothesis: {hypothesis} + \n  Premise: {premise}" ""
+        f"""{instruction}""" + "\n Hypothesis: {hypothesis} + \n  Premise: {premise}" ""
     )
     if not zero_shot:
         try:
             test_data = load_xnli_dataset(dataset_name, lang, split="test", limit=100)
         except KeyError as e:
-            raise KeyError(f"{lang} is not supported in {dataset_name} dataset, choose supported language in few-shot")
+            raise KeyError(
+                f"{lang} is not supported in {dataset_name} dataset, choose supported language in few-shot"
+            )
 
     ic_examples = []
     if not zero_shot:
@@ -425,4 +425,5 @@ def construct_prompt(
         )
 
     return prompt.format(
-        hypothesis=test_example["hypothesis"], premise=test_example["premise"])
+        hypothesis=test_example["hypothesis"], premise=test_example["premise"]
+    )

tasks/qa.py

@@ -10,8 +10,6 @@ import unicodedata
 from typing import Any, Dict, List, NewType, Optional, Union
 
 import numpy as np
-import openai
-import requests
 import yaml
 from datasets import Dataset, load_dataset
 from easygoogletranslate import EasyGoogleTranslate
@@ -20,52 +18,6 @@ from langchain.prompts import FewShotPromptTemplate, PromptTemplate
 from tqdm import tqdm
 from yaml.loader import SafeLoader
 
-
-# from models.model_completion import gpt3x_completion, gemini_completion
-
-def gemini_completion(prompt):
-    # Define the endpoint URL
-    genai.configure(api_key="")
-    model = genai.GenerativeModel("models/gemini-1.0-pro-latest")
-    return model.generate_content(prompt).text
-
-
-
-def get_entities_gpt3_long(prompt):
-    response = openai.ChatCompletion.create(
-        engine="chatgpt", temperature=0, messages=[{"role": "user", "content": prompt}]
-    )
-    return response["choices"][0]["message"]["content"]
-
-
-def gpt3x_completion(
-        prompt: Union[str, List[Dict[str, str]]],
-        model: str = "chatgpt",
-        # run_details: Any = {},
-        # num_evals_per_sec: int = 2,
-        # **model_params,
-) -> str:
-    import openai
-    def get_entities_chatGPT(final_prompt):
-        response = openai.ChatCompletion.create(
-            engine="gpt35-16k",
-            temperature=0,
-            messages=[
-                {"role": "user", "content": final_prompt}
-            ]
-        )
-        return response['choices'][0]['message']['content']
-
-    return get_entities_chatGPT(final_prompt=prompt)
-
-
-def mt0_completion(prompt):
-    inputs = tokenizer.encode(prompt, return_tensors="pt").to("cuda")
-    outputs = model.generate(inputs)
-    return tokenizer.decode(outputs[0])
-
-
-
 XQUAD_LANG2CODES = {
     "bengali": "bn",
     "korean": "ko",
@@ -164,7 +116,7 @@ LANGUAGE_TO_SUFFIX = {
     "hungarian": "hu",
     "czech": "cs",
     "slovak": "sk",
-    "ukrainian": "uk"
+    "ukrainian": "uk",
 }
 
 
@@ -215,20 +167,21 @@ def load_qa_dataset(dataset_name, lang, split, translate_test=False, limit=5):
 
 
 def construct_prompt(
-        instruction: str,
-        test_example: dict,
-        ic_examples: List[dict],
-        zero_shot: bool,
-        lang: str,
-        config: Any,
+    instruction: str,
+    test_example: dict,
+    ic_examples: List[dict],
+    zero_shot: bool,
+    lang: str,
+    config: Any,
 ):
     example_prompt = PromptTemplate(
         input_variables=["context", "question", "answers"],
-        template="Context: {context}  \n  Question: {question} \n " "Answers: {answers}",
+        template="Context: {context}  \n  Question: {question} \n "
+        "Answers: {answers}",
     )
 
     zero_shot_template = (
-            f"""{instruction}""" + " \n <Context>: {context}  \n <Question>: {question} " ""
+        f"""{instruction}""" + " \n <Context>: {context}  \n <Question>: {question} " ""
     )
 
     prompt = (
@@ -260,7 +213,7 @@ def construct_prompt(
 
 
 def dump_metrics(
-        lang: str, config: Dict[str, str], f1: float, em: float, metric_logger_path: str
+    lang: str, config: Dict[str, str], f1: float, em: float, metric_logger_path: str
 ):
     # Check if the metric logger file exists
     file_exists = os.path.exists(metric_logger_path)
@@ -303,7 +256,7 @@ def _translate_instruction(basic_instruction: str, target_language: str) -> str:
 
 
 def _translate_prediction_to_output_language(
-        prediction: str, prediction_language: str, output_language: str
+    prediction: str, prediction_language: str, output_language: str
 ) -> str:
     translator = EasyGoogleTranslate(
         source_language=LANGUAGE_TO_SUFFIX[prediction_language],
@@ -329,7 +282,7 @@ def create_instruction(lang: str, instruction_language: str, expected_output):
 
 
 def _translate_example(
-        example: Dict[str, str], src_language: str, target_language: str
+    example: Dict[str, str], src_language: str, target_language: str
 ):
     translator = EasyGoogleTranslate(
         source_language=LANGUAGE_TO_SUFFIX[src_language],
@@ -340,19 +293,20 @@ def _translate_example(
         return {
             "question": translator.translate(example["question"]),
             "context": translator.translate(example["context"][:2000])
-                       + translator.translate(example["context"][2000:4000])
-                       + translator.translate(example["context"][4000:6000]),
+            + translator.translate(example["context"][2000:4000])
+            + translator.translate(example["context"][4000:6000]),
             "answers": "",
         }
     except Exception as e:
         pass
 
+
 def choose_few_shot_examples(
-        train_dataset: Dataset,
-        few_shot_size: int,
-        context: List[str],
-        selection_criteria: str,
-        lang: str,
+    train_dataset: Dataset,
+    few_shot_size: int,
+    context: List[str],
+    selection_criteria: str,
+    lang: str,
 ) -> List[Dict[str, Union[str, int]]]:
     """Selects few-shot examples from training datasets
 
@@ -423,7 +377,7 @@ def normalize_answer(s):
 
 
 def process_test_example(
-        test_data, config_header, idx, test_example, config, zero_shot, lang, params
+    test_data, config_header, idx, test_example, config, zero_shot, lang, params
 ):
     try:
         # Your existing code for processing each test example
@@ -456,7 +410,9 @@ def process_test_example(
         )
 
         print(len(prompt))
-        pred = get_prediction(prompt=prompt, endpoint_id=7327255438662041600, project_id=16514800572)
+        pred = get_prediction(
+            prompt=prompt, endpoint_id=7327255438662041600, project_id=16514800572
+        )
         # pred = mixtral_completion(prompt)
         print(pred)
 
@@ -551,10 +507,10 @@ def run_one_configuration(params: Optional[PARAMS] = None):
                 response=pred,
                 label=label,
                 response_logger_file=f"{params['response_logger_root']}"
-                                     + f"/{params['model']}"
-                                     + f"/{lang}/"
-                                     + config_header
-                                     + ".csv",
+                + f"/{params['model']}"
+                + f"/{lang}/"
+                + config_header
+                + ".csv",
             )
 
         except Exception as e:
@@ -572,7 +528,6 @@ def run_one_configuration(params: Optional[PARAMS] = None):
     )
 
 
-
 def run_one_configuration_paralle(params: Optional[PARAMS] = None, zero: bool = False):
     if not params:
         params = read_parameters("../../parameters.yaml")
@@ -624,7 +579,6 @@ def run_one_configuration_paralle(params: Optional[PARAMS] = None, zero: bool =
     pool.join()
 
 
-
 def construct_prompt(
     instruction: str,
     test_example: dict,
@@ -632,10 +586,10 @@ def construct_prompt(
     num_examples: int,
     lang: str,
     config: Dict[str, str],
-    dataset_name: str = 'xquad'
+    dataset_name: str = "xquad",
 ):
     if not instruction:
-        instruction = create_instruction(lang, config['prefix'], config['output'])
+        instruction = create_instruction(lang, config["prefix"], config["output"])
 
     example_prompt = PromptTemplate(
         input_variables=["context", "question", "answers"],
@@ -643,15 +597,16 @@ def construct_prompt(
     )
 
     zero_shot_template = (
-            f"""{instruction}""" + " \n <Context>: {context}  \n <Question>: {question} " ""
+        f"""{instruction}""" + " \n <Context>: {context}  \n <Question>: {question} " ""
     )
     if not zero_shot:
         try:
-            test_data = load_qa_dataset(dataset_name = dataset_name, lang=lang, split="test", limit=100)
+            test_data = load_qa_dataset(
+                dataset_name=dataset_name, lang=lang, split="test", limit=100
+            )
         except Exception as e:
             raise KeyError(f"{lang} is not supported in {dataset_name}")
 
-
     ic_examples = []
     if not zero_shot:
 
@@ -677,12 +632,12 @@ def construct_prompt(
         )
     )
     print("lang", lang)
-    print(config["input"] , lang)
+    print(config["input"], lang)
     if config["input"] != lang:
         test_example = _translate_example(
             example=test_example, src_language=lang, target_language=config["input"]
         )
 
     return prompt.format(
-            question=test_example["question"], context=test_example["context"]
-        )
+        question=test_example["question"], context=test_example["context"]
+    )

tasks/summarization.py

@@ -1,9 +1,9 @@
-from typing import List, Dict, Union
+from typing import Dict, List, Union
 
 import numpy as np
 from datasets import Dataset, load_dataset
 from easygoogletranslate import EasyGoogleTranslate
-from langchain.prompts import PromptTemplate, FewShotPromptTemplate
+from langchain.prompts import FewShotPromptTemplate, PromptTemplate
 
 LANGUAGE_TO_SUFFIX = {
     "chinese_simplified": "zh-CN",
@@ -48,12 +48,16 @@ LANGUAGE_TO_SUFFIX = {
     "hungarian": "hu",
     "czech": "cs",
     "slovak": "sk",
-    "ukrainian": "uk"
+    "ukrainian": "uk",
 }
 
 
 def choose_few_shot_examples(
-        train_dataset: Dataset, few_shot_size: int, context: List[str], selection_criteria: str, lang: str,
+    train_dataset: Dataset,
+    few_shot_size: int,
+    context: List[str],
+    selection_criteria: str,
+    lang: str,
 ) -> List[Dict[str, Union[str, int]]]:
     selected_examples = []
 
@@ -67,13 +71,25 @@ def choose_few_shot_examples(
             .tolist()
         )
 
-    ic_examples = [{'text': train_dataset[idx]['text'], 'summary': train_dataset[idx]['summary']} for idx in
-                   example_idxs]
+    ic_examples = [
+        {"text": train_dataset[idx]["text"], "summary": train_dataset[idx]["summary"]}
+        for idx in example_idxs
+    ]
 
     for idx, ic_language in enumerate(context):
-        selected_examples.append(ic_examples[idx]) if ic_language == lang else (
-            selected_examples.append(
-                _translate_example(example=ic_examples[idx], src_language=lang, target_language=ic_language)))
+        (
+            selected_examples.append(ic_examples[idx])
+            if ic_language == lang
+            else (
+                selected_examples.append(
+                    _translate_example(
+                        example=ic_examples[idx],
+                        src_language=lang,
+                        target_language=ic_language,
+                    )
+                )
+            )
+        )
 
     return selected_examples
 
@@ -87,12 +103,16 @@ def _translate_instruction(basic_instruction: str, target_language: str) -> str:
     return translator.translate(basic_instruction)
 
 
-def _translate_example(example: Dict[str, str], src_language: str, target_language: str):
-    translator = EasyGoogleTranslate(source_language=LANGUAGE_TO_SUFFIX[src_language],
-                                     target_language=LANGUAGE_TO_SUFFIX[target_language],
-                                     timeout=30)
+def _translate_example(
+    example: Dict[str, str], src_language: str, target_language: str
+):
+    translator = EasyGoogleTranslate(
+        source_language=LANGUAGE_TO_SUFFIX[src_language],
+        target_language=LANGUAGE_TO_SUFFIX[target_language],
+        timeout=30,
+    )
     try:
-        return {'text': translator.translate(example['text']), 'summary': ''}
+        return {"text": translator.translate(example["text"]), "summary": ""}
     except Exception as e:
         print(e)
 
@@ -117,17 +137,17 @@ def load_xlsum_data(lang, split, limit=5):
 
 
 def construct_prompt(
-        instruction: str,
-        test_example: dict,
-        zero_shot: bool,
-        dataset: str,
-        num_examples: int,
-        lang: str,
-        config: Dict[str, str],
+    instruction: str,
+    test_example: dict,
+    zero_shot: bool,
+    dataset: str,
+    num_examples: int,
+    lang: str,
+    config: Dict[str, str],
 ):
     if not instruction:
         print(lang)
-        instruction = create_instruction(lang, config['prefix'], config['output'])
+        instruction = create_instruction(lang, config["prefix"], config["output"])
 
     example_prompt = PromptTemplate(
         input_variables=["summary", "text"], template="Text: {text}\nSummary: {summary}"
@@ -139,7 +159,9 @@ def construct_prompt(
         try:
             test_data = load_xlsum_data(lang=lang, split="test", limit=100)
         except Exception as e:
-            raise KeyError(f"{lang} is not supported in XlSum dataset, choose supported language in few-shot")
+            raise KeyError(
+                f"{lang} is not supported in XlSum dataset, choose supported language in few-shot"
+            )
 
     ic_examples = []
     if not zero_shot: