Update app.py

app.py

@@ -12,67 +12,89 @@ tokenizer = AutoTokenizer.from_pretrained(model_id)
 model = AutoModelForCausalLM.from_pretrained(
     model_id,
     device_map="auto",
-    torch_dtype=torch.bfloat16  # Usa bf16 como en el ejemplo original
+    torch_dtype=torch.bfloat16
 )
 
-languages = [ 'Aragonese', 'Asturian', 'Basque', 'Bulgarian', 'Catalan', 'Valencian', 'Croatian', 'Czech', 'Danish', 'Dutch', 'English', 'Estonian', 
-             'Finnish', 'French', 'Galician', 'German', 'Greek', 'Hungarian', 'Irish', 'Italian', 'Latvian', 'Lithuanian', 'Maltese', 'Norwegian Bokmål', 
-             'Norwegian Nynorsk', 'Occitan', 'Aranese', 'Polish', 'Portuguese', 'Romanian', 'Russian', 'Serbian', 'Slovak', 'Slovenian', 'Spanish', 'Swedish', 
-             'Ukrainian', 'Welsh' ]
-
-example_sentence = ["Ahir se'n va anar, va agafar les seves coses i es va posar a navegar."]
+languages = sorted([ 'Aragonese', 'Asturian', 'Basque', 'Bulgarian', 'Catalan', 'Valencian', 'Croatian', 'Czech', 'Danish', 'Dutch', 'English', 'Estonian',
+             'Finnish', 'French', 'Galician', 'German', 'Greek', 'Hungarian', 'Irish', 'Italian', 'Latvian', 'Lithuanian', 'Maltese', 'Norwegian Bokmål',
+             'Norwegian Nynorsk', 'Occitan', 'Aranese', 'Polish', 'Portuguese', 'Romanian', 'Russian', 'Serbian', 'Slovak', 'Slovenian', 'Spanish', 'Swedish',
+             'Ukrainian', 'Welsh' ])
 
 @spaces.GPU(duration=120)
-def translate(input_text, source, target):
-    sentences = [s for s in input_text.strip().split('\n') if s.strip()]
-    translated_sentences = []
-
-    for sentence in sentences:
-        prompt_text = f"Translate the following text from {source} into {target}.\n{source}: {sentence} \n{target}:"
-        messages = [{"role": "user", "content": prompt_text}]
-        date_string = datetime.today().strftime('%Y-%m-%d')
-
-        prompt = tokenizer.apply_chat_template(
-            messages,
-            tokenize=False,
-            add_generation_prompt=True,
-            date_string=date_string
-        )
-
-        inputs = tokenizer(prompt, return_tensors="pt", add_special_tokens=False).to(model.device)
-        input_length = inputs.input_ids.shape[1]
-
-        output = model.generate(
-            input_ids=inputs.input_ids,
-            max_new_tokens=400,
-            early_stopping=True,
-            num_beams=5
-        )
-
-        decoded = tokenizer.decode(output[0, input_length:], skip_special_tokens=True).strip()
-        translated_sentences.append(decoded)
-
-    return '\n'.join(translated_sentences), ""
+def generate_output(task, source, target, input_text, mt_text=None):
+    date_string = datetime.today().strftime('%Y-%m-%d')
+
+    if task == "Translation":
+        prompt = f"Translate the following text from {source} into {target}.\n{source}: {input_text.strip()} \n{target}:"
+    elif task == "Post-editing":
+        if not mt_text:
+            return "Please provide machine translation (MT) for post-editing.", ""
+        prompt = f"Please fix any mistakes in the following {source}-{target} machine translation or keep it unedited if it's correct.\nSource: {input_text.strip()} \nMT: {mt_text.strip()} \nCorrected:"
+    elif task == "Document translation":
+        prompt = f"Please translate this text from {source} into {target}.\n{source}: {input_text.strip()}\n{target}:"
+    elif task == "Grammar checker":
+        prompt = f"Please fix any mistakes in the following {source} sentence or keep it unedited if it's correct.\nSentence: {input_text.strip()} \nCorrected:"
+    elif task == "Named-entity recognition":
+        prompt = """Analyse the following tokenized text and mark the tokens containing named entities.
+Use the following annotation guidelines with these tags for named entities:
+- ORG (Refers to named groups or organizations)
+- PER (Refers to individual people or named groups of people)
+- LOC (Refers to physical places or natural landmarks)
+- MISC (Refers to entities that don't fit into standard categories).
+Prepend B- to the first token of a given entity and I- to the remaining ones if they exist.
+If a token is not a named entity, label it as O.
+Input: """ + str(input_text.strip()) + "\nMarked:"
+
+    messages = [{"role": "user", "content": prompt}]
+    final_prompt = tokenizer.apply_chat_template(
+        messages,
+        tokenize=False,
+        add_generation_prompt=True,
+        date_string=date_string
+    )
+
+    inputs = tokenizer(final_prompt, return_tensors="pt", add_special_tokens=False).to(model.device)
+    input_length = inputs.input_ids.shape[1]
+
+    output = model.generate(
+        input_ids=inputs.input_ids,
+        max_new_tokens=512,
+        early_stopping=True,
+        num_beams=5
+    )
+
+    decoded = tokenizer.decode(output[0, input_length:], skip_special_tokens=True).strip()
+    return decoded, ""
 
 with gr.Blocks() as demo:
-    gr.HTML("""<html>
-  <head><style>h1 { text-align: center; }</style></head>
-  <body><h1>SalamandraTA 7B Translate</h1></body>
-</html>""")
+    gr.Markdown("# 🦎 SalamandraTA 7B - Multitask Demo")
+    gr.Markdown("Explore the translation, grammar correction, NER and post-editing capabilities of the SalamandraTA 7B model.")
 
     with gr.Row():
-        with gr.Column():
-            source_language_dropdown = gr.Dropdown(choices=languages, value="Catalan", label="Source Language")
-            input_textbox = gr.Textbox(lines=5, placeholder="Enter text to translate", label="Input Text")
-        with gr.Column():
-            target_language_dropdown = gr.Dropdown(choices=languages, value="English", label="Target Language")
-            translated_textbox = gr.Textbox(lines=5, placeholder="", label="Translated Text")
+        task_selector = gr.Radio(["Translation", "Document translation", "Post-editing", "Grammar checker", "Named-entity recognition"], value="Translation", label="Select Task")
+
+    with gr.Row():
+        source_lang = gr.Dropdown(choices=languages, value="Catalan", label="Source Language")
+        target_lang = gr.Dropdown(choices=languages, value="English", label="Target Language")
+
+    input_textbox = gr.Textbox(lines=6, placeholder="Enter source text or token list here", label="Input Text")
+    mt_textbox = gr.Textbox(lines=4, placeholder="(Only for Post-editing) Enter machine translation", label="Machine Translation (optional)")
+    output_textbox = gr.Textbox(lines=6, label="Output")
 
     info_label = gr.HTML("")
-    btn = gr.Button("Translate")
-    btn.click(translate, inputs=[input_textbox, source_language_dropdown, target_language_dropdown],
-              outputs=[translated_textbox, info_label])
-    gr.Examples(example_sentence, inputs=[input_textbox])
+    translate_btn = gr.Button("Generate")
+    translate_btn.click(generate_output, inputs=[task_selector, source_lang, target_lang, input_textbox, mt_textbox], outputs=[output_textbox, info_label])
+
+    gr.Examples(
+        examples=[
+            ["Translation", "Catalan", "Galician", "Als antics egipcis del període de l'Imperi Nou els fascinaven els monuments dels seus predecessors, que llavors tenien més de mil anys.", ""],
+            ["Post-editing", "Catalan", "English", "Rafael Nadal i Maria Magdalena van inspirar a una generació sencera.", "Rafael Christmas and Maria the Muffin inspired an entire generation each in their own way."],
+            ["Grammar checker", "Catalan", "", "Entonses, el meu jefe m’ha dit que he de treballar els fins de setmana.", ""],
+            ["Named-entity recognition", "", "", "['La', 'defensa', 'del', 'antiguo', 'responsable', 'de', 'la', 'RFEF', 'confirma', 'que', 'interpondrá', 'un', 'recurso.']", ""],
+            ["Document translation", "English", "Asturian", "President Donald Trump, who campaigned on promises to crack down on illegal immigration, has raised alarms in the U.S. dairy industry...", ""]
+        ],
+        inputs=[task_selector, source_lang, target_lang, input_textbox, mt_textbox]
+    )
 
 if __name__ == "__main__":
     demo.launch()