Merge branch 'main' of https://huggingface.co/spaces/ltg/no-en-translation

app.py

@@ -1,7 +1,7 @@
 import torch
-from transformers import AutoTokenizer, TextIteratorStreamer
-# from modeling_nort5 import NorT5ForConditionalGeneration
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, TextIteratorStreamer
 from threading import Thread
+import gradio as gr
 
 
 print(f"Starting to load the model to memory")
@@ -10,9 +10,10 @@ tokenizer = AutoTokenizer.from_pretrained("nort5_en-no_base")
 cls_index = tokenizer.convert_tokens_to_ids("[CLS]")
 sep_index = tokenizer.convert_tokens_to_ids("[SEP]")
 eos_index = tokenizer.convert_tokens_to_ids("[EOS]")
-eng_index = tokenizer.convert_tokens_to_ids(">>ENG<<")
-nob_index = tokenizer.convert_tokens_to_ids(">>NOB<<")
-nno_index = tokenizer.convert_tokens_to_ids(">>NNO<<")
+pad_index = tokenizer.convert_tokens_to_ids("[PAD]")
+eng_index = tokenizer.convert_tokens_to_ids(">>eng<<")
+nob_index = tokenizer.convert_tokens_to_ids(">>nob<<")
+nno_index = tokenizer.convert_tokens_to_ids(">>nno<<")
 
 model = AutoModelForSeq2SeqLM.from_pretrained("nort5_en-no_base", trust_remote_code=True)
 
@@ -25,14 +26,6 @@ model.eval()
 print(f"Sucessfully loaded the model to the memory")
 
 
-INITIAL_PROMPT = "Du er NorT5, en språkmodell laget ved Universitetet i Oslo. Du er en hjelpsom og ufarlig assistent som er glade for å hjelpe brukeren med enhver forespørsel."
-TEMPERATURE = 0.7
-SAMPLE = True
-BEAMS = 1
-PENALTY = 1.2
-TOP_K = 64
-TOP_P = 0.95
-
 LANGUAGES = [
     "🇬🇧 English",
     "🇳🇴 Norwegian (Bokmål)",
@@ -42,10 +35,44 @@ LANGUAGES = [
 LANGUAGE_IDS = {
     "🇬🇧 English": eng_index,
     "🇳🇴 Norwegian (Bokmål)": nob_index,
-    "🇳🇴 Norwegian (Nynorsk)", nno_index
+    "🇳🇴 Norwegian (Nynorsk)": nno_index
 }
 
 
+class BatchStreamer(TextIteratorStreamer):
+    def put(self, value):
+        print(value.shape)
+
+        #if value.size(0) == 1:
+        #    return super().put(value)
+
+        if len(self.token_cache) == 0:
+            self.token_cache = [[] for _ in range(value.size(0))]
+
+        value = value.tolist()
+
+        # Add the new token to the cache and decodes the entire thing.
+        for c, v in zip(self.token_cache, value):
+            c += [v] if isinstance(v, int) else v
+
+        paragraphs = [tokenizer.decode(c, **self.decode_kwargs).strip() for c in self.token_cache]
+        text = '\n'.join(paragraphs)
+
+        self.on_finalized_text(text)
+
+    def end(self):
+        if len(self.token_cache) > 0:
+            paragraphs = [tokenizer.decode(c, **self.decode_kwargs).strip() for c in self.token_cache]
+            printable_text = '\n'.join(paragraphs)
+            self.token_cache = []
+            self.print_len = 0
+        else:
+            printable_text = ""
+
+        self.next_tokens_are_prompt = True
+        self.on_finalized_text(printable_text, stream_end=True)
+
+
 def set_default_target():
     return "*Translating...*"
 
@@ -54,33 +81,45 @@ def translate(source, source_language, target_language):
     if source_language == target_language:
         return source
 
+    source = [s.strip() for s in source.split('\n')]
     source_subwords = tokenizer(source).input_ids
-    source_subwords = [cls_index, LANGUAGE_IDS[target_language], LANGUAGE_IDS[source_language]] + source_subwords + [sep_index]
-    source_subwords = torch.tensor([source_subwords[:512]])
+    source_subwords = [[cls_index, LANGUAGE_IDS[target_language], LANGUAGE_IDS[source_language]] + s + [sep_index] for s in source_subwords]
+    source_subwords = [torch.tensor(s) for s in source_subwords]
+    source_subwords = torch.nn.utils.rnn.pad_sequence(source_subwords, batch_first=True, padding_value=pad_index)
+    source_subwords = source_subwords[:, :512].to(device)
+
+    streamer = BatchStreamer(tokenizer, timeout=60.0, skip_special_tokens=True)
 
-    predictions = model.generate(
+    def generate(model, **kwargs):
+        with torch.inference_mode():
+            with torch.autocast(enabled=device != "cpu", device_type=device, dtype=torch.bfloat16):
+                return model.generate(**kwargs)
+
+    generate_kwargs = dict(
+        streamer=streamer,
         input_ids=source_subwords,
+        attention_mask=(source_subwords != pad_index).long(),
         max_new_tokens = 512-1,
-        do_sample=False
+        # num_beams=4,
+        # early_stopping=True,
+        do_sample=False,
+        use_cache=True
     )
-    predictions = [tokenizer.decode(p, skip_special_tokens=True) for p in predictions.tolist()]
+    t = Thread(target=generate, args=(model,), kwargs=generate_kwargs)
+    t.start()
 
-    return predictions
+    for new_text in streamer:
+        yield new_text.strip()
+    return new_text.strip()
 
 
 def switch_inputs(source, target, source_language, target_language):
     return target, source, target_language, source_language
 
 
-import gradio as gr
-
 with gr.Blocks(theme='sudeepshouche/minimalist') as demo:
 
     gr.Markdown("# Norwegian-English translation")
-    # gr.HTML('<img src="https://huggingface.co/ltg/norbert3-base/resolve/main/norbert.png" width=6.75%>')
-    # gr.Checkbox(label="I want to publish all my conversations", value=True)
-
-    # chatbot = gr.Chatbot(value=[[None, "Hei, hva kan jeg gjøre for deg? 😊"]])
 
     with gr.Row():
         with gr.Column(scale=7, variant="panel"):
@@ -116,8 +155,8 @@ with gr.Blocks(theme='sudeepshouche/minimalist') as demo:
         return {
             source: gr.update(interactive=True),
             submit: gr.update(interactive=True),
-            source_language: gr.update(interactive=False),
-            target_language: gr.update(interactive=False)
+            source_language: gr.update(interactive=True),
+            target_language: gr.update(interactive=True)
         }