Update app.py

app.py

@@ -4,23 +4,25 @@ from transformers import AutoTokenizer, TextIteratorStreamer
 from threading import Thread
 
 
-# print(f"Starting to load the model to memory")
+print(f"Starting to load the model to memory")
 
-# tokenizer = AutoTokenizer.from_pretrained("nort5_en-no_base")
-# cls_index = tokenizer.convert_tokens_to_ids("[CLS]")
-# sep_index = tokenizer.convert_tokens_to_ids("[SEP]")
-# user_index = tokenizer.convert_tokens_to_ids("[USER]")
-# assistent_index = tokenizer.convert_tokens_to_ids("[ASSISTENT]")
+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<<")
 
-# model = NorT5ForConditionalGeneration.from_pretrained("nort5_en-no_base", ignore_mismatched_sizes=True)
+model = AutoModelForSeq2SeqLM.from_pretrained("nort5_en-no_base", trust_remote_code=True)
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 print(f"SYSTEM: Running on {device}", flush=True)
 
-# model = model.to(device)
-# model.eval()
+model = model.to(device)
+model.eval()
 
-# print(f"Sucessfully loaded the model to the memory")
+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."
@@ -37,13 +39,33 @@ LANGUAGES = [
     "🇳🇴 Norwegian (Nynorsk)"
 ]
 
+LANGUAGE_IDS = {
+    "🇬🇧 English": eng_index,
+    "🇳🇴 Norwegian (Bokmål)": nob_index,
+    "🇳🇴 Norwegian (Nynorsk)", nno_index
+}
+
 
 def set_default_target():
     return "*Translating...*"
 
 
 def translate(source, source_language, target_language):
-    return "This is a fake translation"
+    if source_language == target_language:
+        return source
+
+    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]])
+
+    predictions = model.generate(
+        input_ids=source_subwords,
+        max_new_tokens = 512-1,
+        do_sample=False
+    )
+    predictions = [tokenizer.decode(p, skip_special_tokens=True) for p in predictions.tolist()]
+
+    return predictions
 
 
 def switch_inputs(source, target, source_language, target_language):