- Update lemmatisering

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+*.json_large filter=lfs diff=lfs merge=lfs -text

README.md

@@ -20,6 +20,7 @@ This specific version of the tagger is based on Norbert3-base.
 The aim of this model is to make Humit-Tagger available as a HuggingFace model including all functionality that the [original code](https://github.com/humit-oslo/humit-tagger) supports.
 In addition to the morphological tagging, this model supports Nynorsk/Bokmåk language identification provided by this [repository](https://github.com/humit-oslo/humit-sprakidentifikator).
 
+**If you do not need lemmatisation, we recommend giving the lemmatisation=False flag as parameter. This will save some cpu time and make the tagging faster.**
 
 This model adds four classification layers on top of the base model.
 These layers do language identification, morphologic classification, lemmatization classification, and sentence boundary detection.
@@ -75,6 +76,7 @@ These functions receive similar parameters.
 | lang\_per\_sentence| yes | no | True / False | False | identify the language per sentence and output the tags according to the language identified for that sentence. If this is not set, and lang is "au" then the whole input (or a file if input\_directory is used) is used to identify the language. |
 | write\_output\_to | yes | yes | a file path, a file handle, or "list" | sys.stdout | to specify where to write the output. If a file path is provided, the output will be written to that file. The file is overwritten. If a file handle is provided, then the output is written there. If "list" is given as parameters, then the function returns a python "list". |
 | output\_tsv | yes | yes | True/False | False | to specify the output format. The default is the json format. If multiple sentences exist, each line is a single valid json but not the whole output. This option cannot be used along with write\_output\_to="list" |
+| lemmatisation | yes | no | True / False | True | to specify whether lemmatisation will be applied. Disabling lemmatisation by giving this parameter as False makes the tagger faster. |
 | lang\_per\_item | no | yes | True/False | False | consider each item in the list given as separate input for language identification. |
 | fast\_mode | no  | yes | True/False | False | identify languages of the files in the input directory in fast mode. This mode uses only the beginning of the files in identification. This method is much more faster for many files but is not as accurate as if this paramer is set to False. |
 

fullform_list.json_large

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a32e9d7c36ed2ba9ec7f080e118760c444277fb6f213172246d24711b0493433
+size 240703613

modeling_humit_tagger.py

@@ -32,7 +32,7 @@ class HumitTaggerModel(torch.nn.Module):
             kwargs["this_model_config"]=json.load(js)
 
 
-        # Download this model's config:
+        # Download this model's lemma rules pickle file:
         lemma_rules_path = hf_hub_download(repo_id=repo_name, filename=kwargs["config"].lemma_rules_py_file)
 
         # load lemma rules class
@@ -46,6 +46,7 @@ class HumitTaggerModel(torch.nn.Module):
         base_config_file = hf_hub_download(repo_id=kwargs["this_model_config"]["base_model"], filename=kwargs["this_model_config"]["base_model_config_file"])
         base_model_file = hf_hub_download(repo_id=kwargs["this_model_config"]["base_model"], filename=kwargs["this_model_config"]["base_model_model_file"])
         base_model_config_json_file = hf_hub_download(repo_id=kwargs["this_model_config"]["base_model"], filename=kwargs["this_model_config"]["base_model_config_json_file"])
+        fullformlist_file = hf_hub_download(repo_id=repo_name, filename=kwargs["this_model_config"]["fullformlist_file"])
 
         # Copy base model's configuration python file into our working directory
         config_file_path = os.path.join(os.path.dirname(os.path.abspath(__file__)) , os.path.basename(base_config_file))
@@ -81,12 +82,13 @@ class HumitTaggerModel(torch.nn.Module):
 
         kwargs["model_weights_path"] = model_weights_path
         kwargs["repo_name"] = repo_name
+        kwargs["fullformlist_file"] = fullformlist_file
         return HumitTaggerModel(**kwargs)
 
     def __init__(self, **kwargs ): 
         super(HumitTaggerModel, self).__init__()
         json_cfg = kwargs["base_model_json_cfg"]
-        self.config=kwargs["this_model_config"]
+        self.config = kwargs["this_model_config"]
         self.LemmaHandling = sys.modules["lemma_rules"].LemmaHandling
         self.LemmaHandling.load_lemma_rules_from_obj(self.config["lemma_rules"])
         cfg=sys.modules["base_config"].NorbertConfig(**json_cfg)
@@ -117,6 +119,32 @@ class HumitTaggerModel(torch.nn.Module):
         self.REPLACE_PATTERN = '|'.join(sorted(re.escape(k) for k in self.REPLACE_DICT))
         self.MAX_LENGTH = self.bert.config.max_position_embeddings
 
+        # Note the classes that represents gen and prop tags
+        self.gen_tag_classes = set()
+        self.prop_tag_classes = set()
+        self.t_2_tag_classes = set()
+
+        for i, lst in enumerate(self.config["tags"][0]):
+            if "gen" in lst:
+                self.gen_tag_classes.add(i)
+            if "prop" in lst:
+                self.prop_tag_classes.add(i)
+            if "2" in lst:
+                self.t_2_tag_classes.add(i)
+
+
+        # Load the fullform list
+        self.fullform_list=[{},{}]
+        try:
+            with open(kwargs["fullformlist_file"], 'r') as f:
+                self.fullform_list = json.load(f)
+                for k in range(2):
+                    for i in self.fullform_list[k]:
+                        for j in self.fullform_list[k][i][j]:
+                            self.fullform_list[k][i][j]=set(self.fullform_list[k][i][j])
+        except:
+            pass
+
     def forward(self, input_ids=None, attention_mask=None ):
         outputs = self.bert(input_ids=input_ids, attention_mask=attention_mask, return_dict=True )
         sequence_output = self.dropout(outputs.last_hidden_state)
@@ -171,19 +199,24 @@ class HumitTaggerModel(torch.nn.Module):
                             }
             batched_sentences.append(to_append)
 
-        torch.cuda.empty_cache()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
 
         return batched_sentences
 
     def _split_sentences(self, inp):
 
+        # Remove double spaces
+        inp=" ".join(inp.split())
+
         # Here we get the whole text tokenized.
         encodings = self.tokenizer(inp,add_special_tokens=False, return_tensors="pt").to(self.device)
 
         # Save a copy of the tokenization
         original_encodings=copy.deepcopy(encodings)
         original_encodings=original_encodings.to("cpu")
-        torch.cuda.empty_cache()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
 
         # Pad to the complete size (model max_size -1 (-1 to add CLS))
         old_size=encodings["input_ids"][0].size()[0]
@@ -225,13 +258,15 @@ class HumitTaggerModel(torch.nn.Module):
         # First get them back to CPU to open space on GPU
         input_ids_batched=[i.to("cpu") for i in input_ids_batched]
         attention_mask_batched=[i.to("cpu") for i in attention_mask_batched]
-        torch.cuda.empty_cache()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
 
         for input_ids, attention_masks in zip(input_ids_batched, attention_mask_batched):
             current_batch={"input_ids":input_ids.to(self.device).long(), "attention_mask":attention_masks.to(self.device).long()}
             outputs = self(**current_batch)
             del current_batch
-            torch.cuda.empty_cache()
+            if torch.cuda.is_available():
+                torch.cuda.empty_cache()
 
             label_data=outputs["logits1"].argmax(-1)
             labels_output.extend(label_data)
@@ -240,7 +275,8 @@ class HumitTaggerModel(torch.nn.Module):
         labels_output=torch.stack(labels_output ,dim=0)
         labels_output=labels_output[:, range(1,self.MAX_LENGTH)]
         labels_output=torch.reshape(labels_output,(1,row_count * self.MAX_LENGTH_WITHOUT_CLS))
-        torch.cuda.empty_cache()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
 
         # Now the data is split into sentences
         # So, now create sentence data as list so that this could be used
@@ -265,7 +301,9 @@ class HumitTaggerModel(torch.nn.Module):
         del old_size
         del inp
         del outputs
-        torch.cuda.empty_cache()
+
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
 
         return sentence_list
 
@@ -279,6 +317,85 @@ class HumitTaggerModel(torch.nn.Module):
             sentences.extend(self._split_sentences(i.strip()))
         return sentences
 
+    def _lemmatize(self, tag, LANG):
+
+        # Here, a "tag" is a list of words in one sentence, their tags and an ordering of lemma classes according the lemmatization model for each word.
+        # We go over all words, and apply our algorithm for lemmatization
+        # 1. If the "pron" tag is found in the tags
+        #    then, we check if the "gen" tag also exists
+        #         if there is the "gen" tag in tags and if there is "s" at the end of the word, we remove that s
+        #             and return the rest of the word as lemma
+        # 2. OR, we continue with "høflig" processing
+        #    if the word is "De" and if it has the tag "høflig" then we set the lemma as "De", otherwise "de"
+        # 3. OR, we continue with checking the word and its word class (subst, verb, adj, etc.) towards the fullform lists.
+        #    if the word and its word class exists in the fullformlist (of the language bokmål or nynorsk according the the language parameter)
+        #       then we set the lemma from the fullform list.
+        #       if there are multiple lemmas in the fullform list, then we check each lemma suggested by the model
+        #          we pick the lemma amon the lemmas suggested by the fullformlist that comes the first among the lemmas suggested by model
+        # 4. OR, we set the first lemma suggested by the model
+        # 5. OR, just in case, one way or another if we cannot set a lemma, we set the word as the lemma
+
+        # Go over all words in the sentence
+        for i in range(len(tag)):
+
+            # If there is prop in tags
+            if tag[i]["t"] in self.prop_tag_classes:
+
+                # set the lemma as the word
+                tag[i]["l"]=tag[i]["w"]
+
+                # if there is gen in tags then remove the last Ss
+                if tag[i]["t"] in self.gen_tag_classes:
+                    if tag[i]["l"].endswith("'s") or tag[i]["l"].endswith("'S"):
+                        tag[i]["l"]=tag[i]["l"][:-2]
+                    elif tag[i]["l"].endswith("s") or tag[i]["l"].endswith("S") or tag[i]["l"].endswith("'"):
+                        tag[i]["l"]=tag[i]["l"][:-1]
+                continue
+
+            # if høflig
+            if tag[i]["w"]=="De":
+                if tag[i]["t"] in self.t_2_tag_classes:
+                    tag[i]["l"]="De"
+                    continue
+                else:
+                    tag[i]["l"]="de"
+                    continue
+
+            # for the rest of the cases of the word, lowercase the word and check against the fullform list
+            word=tag[i]["w"].lower()
+            word_class = self.tags[0][tag[i]["t"]][0]
+
+            # get the lemma from the fullform list
+            fullform_list_lemma = self.fullform_list[LANG].get(word, {}).get(word_class)
+
+            # if there is not a lemma in the fullformlist
+            # use the first lemma from the model
+            if fullform_list_lemma==None:
+                tag[i]["l"] = self.LemmaHandling.get_lemma_given_word_and_lemma_list_index(tag[i]["w"], tag[i]["l"][0] )
+
+            # if there is only one fullformlist-lemma:
+            elif len(fullform_list_lemma) == 1:
+                tag[i]["l"] = next(iter(fullform_list_lemma))
+
+            # if there are multiple lemmas in the fullformlist
+            # here we disambugate among these lemmas using the alternatives from the model
+            elif len(fullform_list_lemma) > 1:
+                tag[i]["l"] = next((selected_lemma for x in tag[i]["l"] if (selected_lemma :=  self.LemmaHandling.get_lemma_given_word_and_lemma_list_index(tag[i]["w"], x )) in fullform_list_lemma),  self.LemmaHandling.get_lemma_given_word_and_lemma_list_index(tag[i]["w"], tag[i]["l"][0] ) )
+
+            # This branch will probably not be called but kept just in case
+            # If none of the cases above, use the first lemma suggested by the model
+            else:
+                tag[i]["l"] = self.LemmaHandling.get_lemma_given_word_and_lemma_list_index(tag[i]["w"], tag[i]["l"][0] )
+
+            # This if will probable not be true either but kept just in case
+            # If a lemma could not be assigned after all these operations
+            # then asign the word itself
+            # Check by if the lemma field is still a list or if the field-type is string the legth is 0
+            if type(tag[i]["l"]) == list or len(tag[i]["l"]) == 0:
+                tag[i]["l"] = tag[i]["w"]
+
+        return tag
+
     def tag_sentence_list(self, lst, **tag_config):
 
         # If the sentences are not tokenized, tokenize while batching:
@@ -296,62 +413,268 @@ class HumitTaggerModel(torch.nn.Module):
         else:
             tokenized_batches = self._batchify(lst)
         
-        # If language will be identified per sentence
-        if tag_config["lang_per_sentence"]:
-            id_to_lang = self.config["id_to_lang"]
-            # If the output will be to a python list
-            if tag_config["write_output_to"]==None:
-                all_tagged_sentences = []
+        # If lemmatization will be applied
+        if tag_config["lemmatize"]:
+
+            # If language will be identified per sentence
+            if tag_config["lang_per_sentence"]:
+                id_to_lang = self.config["id_to_lang"]
+                # If the output will be to a python list
+                if tag_config["write_output_to"]==None:
+                    all_tagged_sentences = []
+                    for batch in tokenized_batches:
+                        all_out = self(batch["input_ids"].to(self.device), batch["attention_mask"].to(self.device))
+                        batch_tags = torch.argmax(all_out["logits2"], dim=-1)
+                        batch_lemma_indices = torch.topk(all_out["logits3"].flatten(start_dim=2, end_dim=2), len(self.LemmaHandling.lemma_list))
+                        #batch_lemmas = torch.argmax(all_out["logits3"], dim=-1)
+                        batch_langs = torch.argmax(all_out["seq_logits"], dim=-1)
+                        batch["input_ids"].to("cpu")
+                        batch["attention_mask"].to("cpu")
+    
+                        for input_ids, tags, lemmas, lang in zip(batch["input_ids"].tolist(), batch_tags.tolist(),
+                                                                 batch_lemma_indices.indices.tolist(), batch_langs[:, 0].tolist()):
+                            this_sentence=[]
+                            for inps, tag, lemma in zip(input_ids[1:], tags[1:], lemmas[1:]):
+                                if inps == self.tokenizer.sep_token_id or inps == self.tokenizer.pad_token_id:
+                                    break
+                                if lemma[0] == 0: # If there is no lemma here, that means we haven't reached the end of the word
+                                    if len(this_sentence)>0:
+                                        this_sentence[-1]["w"] += self.tokenizer.decode(inps)
+                                    else:
+                                        this_sentence.append({"w": self.tokenizer.decode(inps), "t": tag, "l":lemma})
+                                else:
+                                    this_sentence.append({"w":self.tokenizer.decode(inps).strip(), "t":tag, "l":lemma})
+                            this_sentence = self._lemmatize(this_sentence, lang)
+                            all_tagged_sentences.append({"lang":id_to_lang[lang], "sent": [ {"w":i["w"], "t":self.tags[lang][i["t"]], "l":i["l"]} for i in this_sentence]})
+
+                    return all_tagged_sentences
+
+                # If the output is in TSV format to a pipe (stdout or a file handle)
+                elif tag_config["output_tsv"]:
+                    for batch in tokenized_batches:
+                        all_out = self(batch["input_ids"].to(self.device), batch["attention_mask"].to(self.device))
+                        batch_tags = torch.argmax(all_out["logits2"], dim=-1)
+                        batch_lemma_indices = torch.topk(all_out["logits3"].flatten(start_dim=2, end_dim=2), len(self.LemmaHandling.lemma_list))
+                        #batch_lemmas = torch.argmax(all_out["logits3"], dim=-1)
+                        batch_langs = torch.argmax(all_out["seq_logits"], dim=-1)
+                        batch["input_ids"].to("cpu")
+                        batch["attention_mask"].to("cpu")
+
+                        for input_ids, tags, lemmas, lang in zip(batch["input_ids"].tolist(), batch_tags.tolist(),
+                                                                 batch_lemma_indices.indices.tolist(), batch_langs[:, 0].tolist()):
+                            this_sentence=[]
+                            for inps, tag, lemma in zip(input_ids[1:], tags[1:], lemmas[1:]):
+                                if inps == self.tokenizer.sep_token_id or inps == self.tokenizer.pad_token_id:
+                                    break
+                                if lemma[0] == 0: # If there is no lemma here, that means we haven't reached the end of the word
+                                    if len(this_sentence)>0:
+                                        this_sentence[-1]["w"] += self.tokenizer.decode(inps)
+                                    else:
+                                        this_sentence.append({"w": self.tokenizer.decode(inps), "t": tag, "l":lemma})  
+                                else:
+                                    this_sentence.append({"w":self.tokenizer.decode(inps).strip(), "t":tag, "l":lemma})
+                            this_sentence = self._lemmatize(this_sentence, lang)
+                            this_sentence=[ {"w":i["w"], "t":self.tags_str[lang][i["t"]], "l":i["l"]} for i in this_sentence]
+                            tag_config["write_output_to"].write(id_to_lang[lang])
+                            for lin in this_sentence:
+                                tag_config["write_output_to"].write("\t")
+                                tag_config["write_output_to"].write(lin["w"])
+                                tag_config["write_output_to"].write("\t")
+                                tag_config["write_output_to"].write(lin["l"])
+                                tag_config["write_output_to"].write("\t")
+                                tag_config["write_output_to"].write(lin["t"])
+                                tag_config["write_output_to"].write("\n")
+                            tag_config["write_output_to"].write("\n")
+
+                # If output format will be json to a pipe (stdout or a file handle)
+                else:
+                    for batch in tokenized_batches:
+                        all_out = self(batch["input_ids"].to(self.device), batch["attention_mask"].to(self.device))
+                        batch_tags = torch.argmax(all_out["logits2"], dim=-1)
+                        batch_lemma_indices = torch.topk(all_out["logits3"].flatten(start_dim=2, end_dim=2), len(self.LemmaHandling.lemma_list))
+                        #batch_lemmas = torch.argmax(all_out["logits3"], dim=-1)
+                        batch_langs = torch.argmax(all_out["seq_logits"], dim=-1)
+                        batch["input_ids"].to("cpu")
+                        batch["attention_mask"].to("cpu")
+    
+                        for input_ids, tags, lemmas, lang in zip(batch["input_ids"].tolist(), batch_tags.tolist(),
+                                                                 batch_lemma_indices.indices.tolist(), batch_langs[:, 0].tolist()):
+                            this_sentence=[]
+                            for inps, tag, lemma in zip(input_ids[1:], tags[1:], lemmas[1:]):
+                                if inps == self.tokenizer.sep_token_id or inps == self.tokenizer.pad_token_id:
+                                    break
+                                if lemma[0] == 0: # If there is no lemma here, that means we haven't reached the end of the word
+                                    if len(this_sentence)>0:
+                                        this_sentence[-1]["w"] += self.tokenizer.decode(inps)
+                                    else:
+                                        this_sentence.append({"w": self.tokenizer.decode(inps), "t": tag, "l":lemma})
+                                else:
+                                    this_sentence.append({"w":self.tokenizer.decode(inps).strip(), "t":tag, "l":lemma})
+                            this_sentence = self._lemmatize(this_sentence, lang)
+                            json.dump({"lang":id_to_lang[lang], "sent":[ {"w":i["w"], "t":self.tags[lang][i["t"]], "l":i["l"]} for i in this_sentence]}, tag_config["write_output_to"])
+                            tag_config["write_output_to"].write("\n")
+
+            # If the language is set as parameter
+            elif tag_config["lang"] != -1:
+                LANG = tag_config["lang"]
+                LANG_STR = self.config["id_to_lang"][LANG]
+                # If the output will be to a python list
+                if tag_config["write_output_to"]==None:
+                    all_tagged_sentences = []
+                    for batch in tokenized_batches:
+                        all_out = self(batch["input_ids"].to(self.device), batch["attention_mask"].to(self.device))
+                        batch_tags = torch.argmax(all_out["logits2"], dim=-1)                                             
+                        batch_lemma_indices = torch.topk(all_out["logits3"].flatten(start_dim=2, end_dim=2), len(self.LemmaHandling.lemma_list))
+                        #batch_lemmas = torch.argmax(all_out["logits3"], dim=-1)
+                        batch["input_ids"].to("cpu")
+                        batch["attention_mask"].to("cpu")
+                        for input_ids, tags, lemma_indices in zip(batch["input_ids"].tolist(), batch_tags.tolist(),
+                                                                 batch_lemma_indices.indices.tolist()): #batch_lemmas.tolist(),
+                            this_sentence=[]
+                            for inps, tag, lemma in zip(input_ids[1:], tags[1:], lemma_indices[1:]):
+                                if inps == self.tokenizer.sep_token_id or inps == self.tokenizer.pad_token_id:
+                                    break
+                                if lemma[0] == 0: # If there is no lemma here, that means we haven't reached the end of the word
+                                    if len(this_sentence)>0:
+                                        this_sentence[-1]["w"] += self.tokenizer.decode(inps)
+                                    else:
+                                        this_sentence.append({"w": self.tokenizer.decode(inps), "t": tag, "l":lemma})
+                                else:
+                                    this_sentence.append({"w":self.tokenizer.decode(inps).strip(), "t":tag, "l":lemma})
+
+                            this_sentence = self._lemmatize(this_sentence, LANG)
+                            all_tagged_sentences.append({"lang":LANG_STR, "sent": [ {"w":i["w"], "t":self.tags[LANG][i["t"]], "l":i["l"]} for i in this_sentence]})
+                        
+                    return all_tagged_sentences
+                        
+                # If the output is in TSV format to a pipe (stdout or a file handle)
+                elif tag_config["output_tsv"]:
+                    for batch in tokenized_batches:
+                        all_out = self(batch["input_ids"].to(self.device), batch["attention_mask"].to(self.device))
+                        batch_tags = torch.argmax(all_out["logits2"], dim=-1)
+                        batch_lemma_indices = torch.topk(all_out["logits3"].flatten(start_dim=2, end_dim=2), len(self.LemmaHandling.lemma_list))
+                        #batch_lemmas = torch.argmax(all_out["logits3"], dim=-1)
+                        batch["input_ids"].to("cpu")
+                        batch["attention_mask"].to("cpu")
+
+                        for input_ids, tags, lemmas in zip(batch["input_ids"].tolist(), batch_tags.tolist(),
+                                                                 batch_lemma_indices.indices.tolist()):
+                            this_sentence=[]
+                            for inps, tag, lemma in zip(input_ids[1:], tags[1:], lemmas[1:]):
+                                if inps == self.tokenizer.sep_token_id or inps == self.tokenizer.pad_token_id:
+                                    break
+                                if lemma[0] == 0: # If there is no lemma here, that means we haven't reached the end of the word
+                                    if len(this_sentence)>0:
+                                        this_sentence[-1]["w"] += self.tokenizer.decode(inps)
+                                    else:
+                                        this_sentence.append({"w": self.tokenizer.decode(inps), "t": tag, "l":lemma})  
+                                else:
+                                    this_sentence.append({"w":self.tokenizer.decode(inps).strip(), "t":tag, "l":lemma})
+
+                            this_sentence = self._lemmatize(this_sentence, LANG)
+                            this_sentence=[ {"w":i["w"], "t":self.tags_str[LANG][i["t"]], "l":i["l"]} for i in this_sentence]
+                            tag_config["write_output_to"].write(LANG_STR)
+                            for lin in this_sentence:
+                                tag_config["write_output_to"].write("\t")
+                                tag_config["write_output_to"].write(lin["w"])
+                                tag_config["write_output_to"].write("\t")
+                                tag_config["write_output_to"].write(lin["l"])
+                                tag_config["write_output_to"].write("\t")
+                                tag_config["write_output_to"].write(lin["t"])
+                                tag_config["write_output_to"].write("\n")
+                            tag_config["write_output_to"].write("\n")
+
+                # If output format will be json to a pipe (stdout or a file handle)
+                else:
+                    for batch in tokenized_batches:
+                        all_out = self(batch["input_ids"].to(self.device), batch["attention_mask"].to(self.device))
+                        batch_tags = torch.argmax(all_out["logits2"], dim=-1)
+                        batch_lemma_indices = torch.topk(all_out["logits3"].flatten(start_dim=2, end_dim=2), len(self.LemmaHandling.lemma_list))
+                        #batch_lemmas = torch.argmax(all_out["logits3"], dim=-1)
+                        batch["input_ids"].to("cpu")
+                        batch["attention_mask"].to("cpu")
+                        
+                        for input_ids, tags, lemmas in zip(batch["input_ids"].tolist(), batch_tags.tolist(),
+                                                                 batch_lemma_indices.indices.tolist()):
+                            this_sentence=[]
+                            for inps, tag, lemma in zip(input_ids[1:], tags[1:], lemmas[1:]):
+                                if inps == self.tokenizer.sep_token_id or inps == self.tokenizer.pad_token_id:
+                                    break
+                                if lemma[0] == 0: # If there is no lemma here, that means we haven't reached the end of the word
+                                    if len(this_sentence)>0:
+                                        this_sentence[-1]["w"] += self.tokenizer.decode(inps)
+                                    else:
+                                        this_sentence.append({"w": self.tokenizer.decode(inps), "t": tag, "l":lemma})
+                                else:
+                                    this_sentence.append({"w":self.tokenizer.decode(inps).strip(), "t":tag, "l":lemma})
+
+                            this_sentence = self._lemmatize(this_sentence, LANG)
+                            json.dump({"lang":LANG_STR, "sent": [ {"w":i["w"], "t":self.tags[LANG][i["t"]], "l":i["l"]} for i in this_sentence]}, tag_config["write_output_to"])
+                            tag_config["write_output_to"].write("\n")
+
+            # If language will be identified according to the majority of all sentences:
+            else:
+                all_tags=[]
+                all_lemmas=[]
+                all_langs=[]
+                all_input_ids=[]
+                # Go over all batches and each sentence in each batch
                 for batch in tokenized_batches:
                     all_out = self(batch["input_ids"].to(self.device), batch["attention_mask"].to(self.device))
                     batch_tags = torch.argmax(all_out["logits2"], dim=-1)
-                    batch_lemmas = torch.argmax(all_out["logits3"], dim=-1)
+                    batch_lemma_indices = torch.topk(all_out["logits3"].flatten(start_dim=2, end_dim=2), len(self.LemmaHandling.lemma_list))
+                    #batch_lemmas = torch.argmax(all_out["logits3"], dim=-1)
                     batch_langs = torch.argmax(all_out["seq_logits"], dim=-1)
+                    all_input_ids.extend(batch["input_ids"].tolist())
                     batch["input_ids"].to("cpu")
                     batch["attention_mask"].to("cpu")
+                    all_langs.extend(batch_langs[:, 0].tolist())
+                    all_tags.extend(batch_tags.tolist())
+                    all_lemmas.extend(batch_lemma_indices.indices.tolist())
 
-                    for input_ids, tags, lemmas, lang in zip(batch["input_ids"].tolist(), batch_tags.tolist(),
-                                                             batch_lemmas.tolist(), batch_langs[:, 0].tolist()):
+                # Identify the language
+                tag_config["lang"] = 1 if sum(all_langs)/len(all_langs)>=0.5 else 0
+                LANG = tag_config["lang"]
+                LANG_STR = self.config["id_to_lang"][LANG]
+
+                # If the output will be returned as python list:
+                if tag_config["write_output_to"]==None:
+                    all_tagged_sentences = []
+                    for input_ids, tags, lemmas in zip(all_input_ids, all_tags, all_lemmas):
                         this_sentence=[]
                         for inps, tag, lemma in zip(input_ids[1:], tags[1:], lemmas[1:]):
                             if inps == self.tokenizer.sep_token_id or inps == self.tokenizer.pad_token_id:
                                 break
-                            if lemma == 0: # If there is no lemma here, that means we haven't reached the end of the word
+                            if lemma[0] == 0: # If there is no lemma here, that means we haven't reached the end of the word
                                 if len(this_sentence)>0:
                                     this_sentence[-1]["w"] += self.tokenizer.decode(inps)
                                 else:
                                     this_sentence.append({"w": self.tokenizer.decode(inps), "t": tag, "l":lemma})
                             else:
                                 this_sentence.append({"w":self.tokenizer.decode(inps).strip(), "t":tag, "l":lemma})
-                        all_tagged_sentences.append({"lang":id_to_lang[lang], "sent": [ {"w":i["w"], "t":self.tags[lang][i["t"]], "l":self.LemmaHandling.get_lemma_given_word_and_lemma_list_index(i["w"],i["l"])} for i in this_sentence]})
 
-                return all_tagged_sentences
-
-            # If the output is in TSV format to a pipe (stdout or a file handle)
-            elif tag_config["output_tsv"]:
-                for batch in tokenized_batches:
-                    all_out = self(batch["input_ids"].to(self.device), batch["attention_mask"].to(self.device))
-                    batch_tags = torch.argmax(all_out["logits2"], dim=-1)
-                    batch_lemmas = torch.argmax(all_out["logits3"], dim=-1)
-                    batch_langs = torch.argmax(all_out["seq_logits"], dim=-1)
-                    batch["input_ids"].to("cpu")
-                    batch["attention_mask"].to("cpu")
+                        this_sentence = self._lemmatize(this_sentence, LANG)
+                        all_tagged_sentences.append({"lang":LANG_STR, "sent":  [ {"w":i["w"], "t":self.tags[LANG][i["t"]], "l":i["l"]} for i in this_sentence] })
+                    return all_tagged_sentences
 
-                    for input_ids, tags, lemmas, lang in zip(batch["input_ids"].tolist(), batch_tags.tolist(),
-                                                             batch_lemmas.tolist(), batch_langs[:, 0].tolist()):
+                # If the output is in TSV format
+                elif tag_config["output_tsv"]:
+                    for input_ids, tags, lemmas in zip(all_input_ids, all_tags, all_lemmas):
                         this_sentence=[]
                         for inps, tag, lemma in zip(input_ids[1:], tags[1:], lemmas[1:]):
                             if inps == self.tokenizer.sep_token_id or inps == self.tokenizer.pad_token_id:
                                 break
-                            if lemma == 0: # If there is no lemma here, that means we haven't reached the end of the word
+                            if lemma[0] == 0: # If there is no lemma here, that means we haven't reached the end of the word
                                 if len(this_sentence)>0:
                                     this_sentence[-1]["w"] += self.tokenizer.decode(inps)
                                 else:
-                                    this_sentence.append({"w": self.tokenizer.decode(inps), "t": tag, "l":lemma})  
+                                    this_sentence.append({"w": self.tokenizer.decode(inps), "t": tag, "l":lemma})
                             else:
                                 this_sentence.append({"w":self.tokenizer.decode(inps).strip(), "t":tag, "l":lemma})
-                        this_sentence=[ {"w":i["w"], "t":self.tags_str[lang][i["t"]], "l":self.LemmaHandling.get_lemma_given_word_and_lemma_list_index(i["w"],i["l"])} for i in this_sentence]
-                        tag_config["write_output_to"].write(id_to_lang[lang])
+
+                        this_sentence = self._lemmatize(this_sentence, LANG)
+                        this_sentence=[ {"w":i["w"], "t":self.tags_str[LANG][i["t"]], "l":i["l"]} for i in this_sentence]
+                        tag_config["write_output_to"].write(LANG_STR)
                         for lin in this_sentence:
                             tag_config["write_output_to"].write("\t")
                             tag_config["write_output_to"].write(lin["w"])
@@ -362,49 +685,235 @@ class HumitTaggerModel(torch.nn.Module):
                             tag_config["write_output_to"].write("\n")
                         tag_config["write_output_to"].write("\n")
 
-            # If output format will be json to a pipe (stdout or a file handle)
-            else:
-                for batch in tokenized_batches:
-                    all_out = self(batch["input_ids"].to(self.device), batch["attention_mask"].to(self.device))
-                    batch_tags = torch.argmax(all_out["logits2"], dim=-1)
-                    batch_lemmas = torch.argmax(all_out["logits3"], dim=-1)
-                    batch_langs = torch.argmax(all_out["seq_logits"], dim=-1)
-                    batch["input_ids"].to("cpu")
-                    batch["attention_mask"].to("cpu")
-
-                    for input_ids, tags, lemmas, lang in zip(batch["input_ids"].tolist(), batch_tags.tolist(),
-                                                             batch_lemmas.tolist(), batch_langs[:, 0].tolist()):
+                # If output format will be json
+                else:
+                    for input_ids, tags, lemmas in zip(all_input_ids, all_tags, all_lemmas):
                         this_sentence=[]
                         for inps, tag, lemma in zip(input_ids[1:], tags[1:], lemmas[1:]):
                             if inps == self.tokenizer.sep_token_id or inps == self.tokenizer.pad_token_id:
                                 break
-                            if lemma == 0: # If there is no lemma here, that means we haven't reached the end of the word
+                            if lemma[0] == 0: # If there is no lemma here, that means we haven't reached the end of the word
                                 if len(this_sentence)>0:
                                     this_sentence[-1]["w"] += self.tokenizer.decode(inps)
                                 else:
                                     this_sentence.append({"w": self.tokenizer.decode(inps), "t": tag, "l":lemma})
                             else:
                                 this_sentence.append({"w":self.tokenizer.decode(inps).strip(), "t":tag, "l":lemma})
-                        
-                        json.dump({"lang":id_to_lang[lang], "sent":[ {"w":i["w"], "t":self.tags[lang][i["t"]], "l":self.LemmaHandling.get_lemma_given_word_and_lemma_list_index(i["w"],i["l"])} for i in this_sentence]}, tag_config["write_output_to"])
+
+                        this_sentence = self._lemmatize(this_sentence, LANG)
+                        json.dump({"lang":LANG_STR, "sent":[ {"w":i["w"], "t":self.tags[LANG][i["t"]], "l":i["l"]} for i in this_sentence]}, tag_config["write_output_to"])
                         tag_config["write_output_to"].write("\n")
 
-        # If the language is set as parameter
-        elif tag_config["lang"] != -1:
-            LANG = tag_config["lang"]
-            LANG_STR = self.config["id_to_lang"][LANG]
-            # If the output will be to a python list
-            if tag_config["write_output_to"]==None:
-                all_tagged_sentences = []
+        # If lemmatization will not be applied:
+        else:
+            # If language will be identified per sentence
+            if tag_config["lang_per_sentence"]:
+                id_to_lang = self.config["id_to_lang"]
+                # If the output will be to a python list
+                if tag_config["write_output_to"]==None:
+                    all_tagged_sentences = []
+                    for batch in tokenized_batches:
+                        all_out = self(batch["input_ids"].to(self.device), batch["attention_mask"].to(self.device))
+                        batch_tags = torch.argmax(all_out["logits2"], dim=-1)
+                        batch_lemmas = torch.argmax(all_out["logits3"], dim=-1)
+                        batch_langs = torch.argmax(all_out["seq_logits"], dim=-1)
+                        batch["input_ids"].to("cpu")
+                        batch["attention_mask"].to("cpu")
+    
+                        for input_ids, tags, lemmas, lang in zip(batch["input_ids"].tolist(), batch_tags.tolist(),
+                                                                 batch_lemmas.tolist(), batch_langs[:, 0].tolist()):
+                            this_sentence=[]
+                            for inps, tag, lemma in zip(input_ids[1:], tags[1:], lemmas[1:]):
+                                if inps == self.tokenizer.sep_token_id or inps == self.tokenizer.pad_token_id:
+                                    break
+                                if lemma == 0: # If there is no lemma here, that means we haven't reached the end of the word
+                                    if len(this_sentence)>0:
+                                        this_sentence[-1]["w"] += self.tokenizer.decode(inps)
+                                    else:
+                                        this_sentence.append({"w": self.tokenizer.decode(inps), "t": tag})
+                                else:
+                                    this_sentence.append({"w":self.tokenizer.decode(inps).strip(), "t":tag})
+                            all_tagged_sentences.append({"lang":id_to_lang[lang], "sent": [ {"w":i["w"], "t":self.tags[lang][i["t"]]} for i in this_sentence]})
+
+                    return all_tagged_sentences
+
+                # If the output is in TSV format to a pipe (stdout or a file handle)
+                elif tag_config["output_tsv"]:
+                    for batch in tokenized_batches:
+                        all_out = self(batch["input_ids"].to(self.device), batch["attention_mask"].to(self.device))
+                        batch_tags = torch.argmax(all_out["logits2"], dim=-1)
+                        batch_lemmas = torch.argmax(all_out["logits3"], dim=-1)
+                        batch_langs = torch.argmax(all_out["seq_logits"], dim=-1)
+                        batch["input_ids"].to("cpu")
+                        batch["attention_mask"].to("cpu")
+
+                        for input_ids, tags, lemmas, lang in zip(batch["input_ids"].tolist(), batch_tags.tolist(),
+                                                                 batch_lemmas.tolist(), batch_langs[:, 0].tolist()):
+                            this_sentence=[]
+                            for inps, tag, lemma in zip(input_ids[1:], tags[1:], lemmas[1:]):
+                                if inps == self.tokenizer.sep_token_id or inps == self.tokenizer.pad_token_id:
+                                    break
+                                if lemma == 0: # If there is no lemma here, that means we haven't reached the end of the word
+                                    if len(this_sentence)>0:
+                                        this_sentence[-1]["w"] += self.tokenizer.decode(inps)
+                                    else:
+                                        this_sentence.append({"w": self.tokenizer.decode(inps), "t": tag})  
+                                else:
+                                    this_sentence.append({"w":self.tokenizer.decode(inps).strip(), "t":tag})
+                            this_sentence=[ {"w":i["w"], "t":self.tags_str[lang][i["t"]] } for i in this_sentence]
+                            tag_config["write_output_to"].write(id_to_lang[lang])
+                            for lin in this_sentence:
+                                tag_config["write_output_to"].write("\t")
+                                tag_config["write_output_to"].write(lin["w"])
+                                tag_config["write_output_to"].write("\t")
+                                tag_config["write_output_to"].write(lin["t"])
+                                tag_config["write_output_to"].write("\n")
+                            tag_config["write_output_to"].write("\n")
+
+                # If output format will be json to a pipe (stdout or a file handle)
+                else:
+                    for batch in tokenized_batches:
+                        all_out = self(batch["input_ids"].to(self.device), batch["attention_mask"].to(self.device))
+                        batch_tags = torch.argmax(all_out["logits2"], dim=-1)
+                        batch_lemmas = torch.argmax(all_out["logits3"], dim=-1)
+                        batch_langs = torch.argmax(all_out["seq_logits"], dim=-1)
+                        batch["input_ids"].to("cpu")
+                        batch["attention_mask"].to("cpu")
+    
+                        for input_ids, tags, lemmas, lang in zip(batch["input_ids"].tolist(), batch_tags.tolist(),
+                                                                 batch_lemmas.tolist(), batch_langs[:, 0].tolist()):
+                            this_sentence=[]
+                            for inps, tag, lemma in zip(input_ids[1:], tags[1:], lemmas[1:]):
+                                if inps == self.tokenizer.sep_token_id or inps == self.tokenizer.pad_token_id:
+                                    break
+                                if lemma == 0: # If there is no lemma here, that means we haven't reached the end of the word
+                                    if len(this_sentence)>0:
+                                        this_sentence[-1]["w"] += self.tokenizer.decode(inps)
+                                    else:
+                                        this_sentence.append({"w": self.tokenizer.decode(inps), "t": tag})
+                                else:
+                                    this_sentence.append({"w":self.tokenizer.decode(inps).strip(), "t":tag})
+                            
+                            json.dump({"lang":id_to_lang[lang], "sent":[ {"w":i["w"], "t":self.tags[lang][i["t"]]} for i in this_sentence]}, tag_config["write_output_to"])
+                            tag_config["write_output_to"].write("\n")
+
+            # If the language is set as parameter
+            elif tag_config["lang"] != -1:
+                LANG = tag_config["lang"]
+                LANG_STR = self.config["id_to_lang"][LANG]
+                # If the output will be to a python list
+                if tag_config["write_output_to"]==None:
+                    all_tagged_sentences = []
+                    for batch in tokenized_batches:
+                        all_out = self(batch["input_ids"].to(self.device), batch["attention_mask"].to(self.device))
+                        batch_tags = torch.argmax(all_out["logits2"], dim=-1)                                              
+                        batch_lemmas = torch.argmax(all_out["logits3"], dim=-1)
+                        batch["input_ids"].to("cpu")
+                        batch["attention_mask"].to("cpu")
+                                
+                        for input_ids, tags, lemmas in zip(batch["input_ids"].tolist(), batch_tags.tolist(),
+                                                                 batch_lemmas.tolist()):
+                            this_sentence=[]
+                            for inps, tag, lemma in zip(input_ids[1:], tags[1:], lemmas[1:]):
+                                if inps == self.tokenizer.sep_token_id or inps == self.tokenizer.pad_token_id:
+                                    break
+                                if lemma == 0: # If there is no lemma here, that means we haven't reached the end of the word
+                                    if len(this_sentence)>0:
+                                        this_sentence[-1]["w"] += self.tokenizer.decode(inps)
+                                    else:
+                                        this_sentence.append({"w": self.tokenizer.decode(inps), "t": tag})
+                                else:
+                                    this_sentence.append({"w":self.tokenizer.decode(inps).strip(), "t":tag})
+                            all_tagged_sentences.append({"lang":LANG_STR, "sent": [ {"w":i["w"], "t":self.tags[LANG][i["t"]]} for i in this_sentence]})
+                        
+                    return all_tagged_sentences
+                        
+                # If the output is in TSV format to a pipe (stdout or a file handle)
+                elif tag_config["output_tsv"]:
+                    for batch in tokenized_batches:
+                        all_out = self(batch["input_ids"].to(self.device), batch["attention_mask"].to(self.device))
+                        batch_tags = torch.argmax(all_out["logits2"], dim=-1)
+                        batch_lemmas = torch.argmax(all_out["logits3"], dim=-1)
+                        batch["input_ids"].to("cpu")
+                        batch["attention_mask"].to("cpu")
+
+                        for input_ids, tags, lemmas in zip(batch["input_ids"].tolist(), batch_tags.tolist(),
+                                                                 batch_lemmas.tolist()):
+                            this_sentence=[]
+                            for inps, tag, lemma in zip(input_ids[1:], tags[1:], lemmas[1:]):
+                                if inps == self.tokenizer.sep_token_id or inps == self.tokenizer.pad_token_id:
+                                    break
+                                if lemma == 0: # If there is no lemma here, that means we haven't reached the end of the word
+                                    if len(this_sentence)>0:
+                                        this_sentence[-1]["w"] += self.tokenizer.decode(inps)
+                                    else:
+                                        this_sentence.append({"w": self.tokenizer.decode(inps), "t": tag})  
+                                else:
+                                    this_sentence.append({"w":self.tokenizer.decode(inps).strip(), "t":tag})
+                            this_sentence=[ {"w":i["w"], "t":self.tags_str[LANG][i["t"]]} for i in this_sentence]
+                            tag_config["write_output_to"].write(LANG_STR)
+                            for lin in this_sentence:
+                                tag_config["write_output_to"].write("\t")
+                                tag_config["write_output_to"].write(lin["w"])
+                                tag_config["write_output_to"].write("\t")
+                                tag_config["write_output_to"].write(lin["t"])
+                                tag_config["write_output_to"].write("\n")
+                            tag_config["write_output_to"].write("\n")
+
+                # If output format will be json to a pipe (stdout or a file handle)
+                else:
+                    for batch in tokenized_batches:
+                        all_out = self(batch["input_ids"].to(self.device), batch["attention_mask"].to(self.device))
+                        batch_tags = torch.argmax(all_out["logits2"], dim=-1)
+                        batch_lemmas = torch.argmax(all_out["logits3"], dim=-1)
+                        batch["input_ids"].to("cpu")
+                        batch["attention_mask"].to("cpu")
+                        
+                        for input_ids, tags, lemmas in zip(batch["input_ids"].tolist(), batch_tags.tolist(),
+                                                                 batch_lemmas.tolist()):
+                            this_sentence=[]
+                            for inps, tag, lemma in zip(input_ids[1:], tags[1:], lemmas[1:]):
+                                if inps == self.tokenizer.sep_token_id or inps == self.tokenizer.pad_token_id:
+                                    break
+                                if lemma == 0: # If there is no lemma here, that means we haven't reached the end of the word
+                                    if len(this_sentence)>0:
+                                        this_sentence[-1]["w"] += self.tokenizer.decode(inps)
+                                    else:
+                                        this_sentence.append({"w": self.tokenizer.decode(inps), "t": tag})
+                                else:
+                                    this_sentence.append({"w":self.tokenizer.decode(inps).strip(), "t":tag})
+                        
+                            json.dump({"lang":LANG_STR, "sent": [ {"w":i["w"], "t":self.tags[LANG][i["t"]]} for i in this_sentence]}, tag_config["write_output_to"])
+                            tag_config["write_output_to"].write("\n")
+
+            # If language will be identified according to the majority of all sentences:
+            else:
+                all_tags=[]
+                all_lemmas=[]
+                all_langs=[]
+                all_input_ids=[]
+                # Go over all batches and each sentence in each batch
                 for batch in tokenized_batches:
                     all_out = self(batch["input_ids"].to(self.device), batch["attention_mask"].to(self.device))
-                    batch_tags = torch.argmax(all_out["logits2"], dim=-1)                                              
+                    batch_tags = torch.argmax(all_out["logits2"], dim=-1)
                     batch_lemmas = torch.argmax(all_out["logits3"], dim=-1)
+                    batch_langs = torch.argmax(all_out["seq_logits"], dim=-1)
+                    all_input_ids.extend(batch["input_ids"].tolist())
                     batch["input_ids"].to("cpu")
                     batch["attention_mask"].to("cpu")
-                                
-                    for input_ids, tags, lemmas in zip(batch["input_ids"].tolist(), batch_tags.tolist(),
-                                                             batch_lemmas.tolist()):
+                    all_langs.extend(batch_langs[:, 0].tolist())
+                    all_tags.extend(batch_tags.tolist())
+                    all_lemmas.extend(batch_lemmas.tolist())
+
+                # Identify the language
+                tag_config["lang"] = 1 if sum(all_langs)/len(all_langs)>=0.5 else 0
+                LANG = tag_config["lang"]
+                LANG_STR = self.config["id_to_lang"][LANG]
+
+                # If the output will be returned as python list:
+                if tag_config["write_output_to"]==None:
+                    all_tagged_sentences = []
+                    for input_ids, tags, lemmas in zip(all_input_ids, all_tags, all_lemmas):
                         this_sentence=[]
                         for inps, tag, lemma in zip(input_ids[1:], tags[1:], lemmas[1:]):
                             if inps == self.tokenizer.sep_token_id or inps == self.tokenizer.pad_token_id:
@@ -413,24 +922,15 @@ class HumitTaggerModel(torch.nn.Module):
                                 if len(this_sentence)>0:
                                     this_sentence[-1]["w"] += self.tokenizer.decode(inps)
                                 else:
-                                    this_sentence.append({"w": self.tokenizer.decode(inps), "t": tag, "l":lemma})
+                                    this_sentence.append({"w": self.tokenizer.decode(inps), "t": tag})
                             else:
-                                this_sentence.append({"w":self.tokenizer.decode(inps).strip(), "t":tag, "l":lemma})
-                        all_tagged_sentences.append({"lang":LANG_STR, "sent": [ {"w":i["w"], "t":self.tags[LANG][i["t"]], "l":self.LemmaHandling.get_lemma_given_word_and_lemma_list_index(i["w"],i["l"])} for i in this_sentence]})
-                        
-                return all_tagged_sentences
-                        
-            # If the output is in TSV format to a pipe (stdout or a file handle)
-            elif tag_config["output_tsv"]:
-                for batch in tokenized_batches:
-                    all_out = self(batch["input_ids"].to(self.device), batch["attention_mask"].to(self.device))
-                    batch_tags = torch.argmax(all_out["logits2"], dim=-1)
-                    batch_lemmas = torch.argmax(all_out["logits3"], dim=-1)
-                    batch["input_ids"].to("cpu")
-                    batch["attention_mask"].to("cpu")
+                                this_sentence.append({"w":self.tokenizer.decode(inps).strip(), "t":tag})                                                                   
+                        all_tagged_sentences.append({"lang":LANG_STR, "sent":  [ {"w":i["w"], "t":self.tags[LANG][i["t"]]} for i in this_sentence] })
+                    return all_tagged_sentences
 
-                    for input_ids, tags, lemmas in zip(batch["input_ids"].tolist(), batch_tags.tolist(),
-                                                             batch_lemmas.tolist()):
+                # If the output is in TSV format
+                elif tag_config["output_tsv"]:
+                    for input_ids, tags, lemmas in zip(all_input_ids, all_tags, all_lemmas):
                         this_sentence=[]
                         for inps, tag, lemma in zip(input_ids[1:], tags[1:], lemmas[1:]):
                             if inps == self.tokenizer.sep_token_id or inps == self.tokenizer.pad_token_id:
@@ -439,32 +939,22 @@ class HumitTaggerModel(torch.nn.Module):
                                 if len(this_sentence)>0:
                                     this_sentence[-1]["w"] += self.tokenizer.decode(inps)
                                 else:
-                                    this_sentence.append({"w": self.tokenizer.decode(inps), "t": tag, "l":lemma})  
+                                    this_sentence.append({"w": self.tokenizer.decode(inps), "t": tag})
                             else:
-                                this_sentence.append({"w":self.tokenizer.decode(inps).strip(), "t":tag, "l":lemma})
-                        this_sentence=[ {"w":i["w"], "t":self.tags_str[LANG][i["t"]], "l":self.LemmaHandling.get_lemma_given_word_and_lemma_list_index(i["w"],i["l"])} for i in this_sentence]
+                                this_sentence.append({"w":self.tokenizer.decode(inps).strip(), "t":tag})
+                        this_sentence=[ {"w":i["w"], "t":self.tags_str[LANG][i["t"]]} for i in this_sentence]
                         tag_config["write_output_to"].write(LANG_STR)
                         for lin in this_sentence:
                             tag_config["write_output_to"].write("\t")
                             tag_config["write_output_to"].write(lin["w"])
                             tag_config["write_output_to"].write("\t")
-                            tag_config["write_output_to"].write(lin["l"])
-                            tag_config["write_output_to"].write("\t")
                             tag_config["write_output_to"].write(lin["t"])
                             tag_config["write_output_to"].write("\n")
                         tag_config["write_output_to"].write("\n")
 
-            # If output format will be json to a pipe (stdout or a file handle)
-            else:
-                for batch in tokenized_batches:
-                    all_out = self(batch["input_ids"].to(self.device), batch["attention_mask"].to(self.device))
-                    batch_tags = torch.argmax(all_out["logits2"], dim=-1)
-                    batch_lemmas = torch.argmax(all_out["logits3"], dim=-1)
-                    batch["input_ids"].to("cpu")
-                    batch["attention_mask"].to("cpu")
-                        
-                    for input_ids, tags, lemmas in zip(batch["input_ids"].tolist(), batch_tags.tolist(),
-                                                             batch_lemmas.tolist()):
+                # If output format will be json
+                else:
+                    for input_ids, tags, lemmas in zip(all_input_ids, all_tags, all_lemmas):
                         this_sentence=[]
                         for inps, tag, lemma in zip(input_ids[1:], tags[1:], lemmas[1:]):
                             if inps == self.tokenizer.sep_token_id or inps == self.tokenizer.pad_token_id:
@@ -473,98 +963,13 @@ class HumitTaggerModel(torch.nn.Module):
                                 if len(this_sentence)>0:
                                     this_sentence[-1]["w"] += self.tokenizer.decode(inps)
                                 else:
-                                    this_sentence.append({"w": self.tokenizer.decode(inps), "t": tag, "l":lemma})
-                            else:
-                                this_sentence.append({"w":self.tokenizer.decode(inps).strip(), "t":tag, "l":lemma})
-                        
-                        json.dump({"lang":LANG_STR, "sent": [ {"w":i["w"], "t":self.tags[LANG][i["t"]], "l":self.LemmaHandling.get_lemma_given_word_and_lemma_list_index(i["w"],i["l"])} for i in this_sentence]}, tag_config["write_output_to"])
-                        tag_config["write_output_to"].write("\n")
-
-        # If language will be identified according to the majority of all sentences:
-        else:
-            all_tags=[]
-            all_lemmas=[]
-            all_langs=[]
-            all_input_ids=[]
-            # Go over all batches and each sentence in each batch
-            for batch in tokenized_batches:
-                all_out = self(batch["input_ids"].to(self.device), batch["attention_mask"].to(self.device))
-                batch_tags = torch.argmax(all_out["logits2"], dim=-1)
-                batch_lemmas = torch.argmax(all_out["logits3"], dim=-1)
-                batch_langs = torch.argmax(all_out["seq_logits"], dim=-1)
-                all_input_ids.extend(batch["input_ids"].tolist())
-                batch["input_ids"].to("cpu")
-                batch["attention_mask"].to("cpu")
-                all_langs.extend(batch_langs[:, 0].tolist())
-                all_tags.extend(batch_tags.tolist())
-                all_lemmas.extend(batch_lemmas.tolist())
-
-            # Identify the language
-            tag_config["lang"] = 1 if sum(all_langs)/len(all_langs)>=0.5 else 0
-            LANG = tag_config["lang"]
-            LANG_STR = self.config["id_to_lang"][LANG]
-
-            # If the output will be returned as python list:
-            if tag_config["write_output_to"]==None:
-                all_tagged_sentences = []
-                for input_ids, tags, lemmas in zip(all_input_ids, all_tags, all_lemmas):
-                    this_sentence=[]
-                    for inps, tag, lemma in zip(input_ids[1:], tags[1:], lemmas[1:]):
-                        if inps == self.tokenizer.sep_token_id or inps == self.tokenizer.pad_token_id:
-                            break
-                        if lemma == 0: # If there is no lemma here, that means we haven't reached the end of the word
-                            if len(this_sentence)>0:
-                                this_sentence[-1]["w"] += self.tokenizer.decode(inps)
+                                    this_sentence.append({"w": self.tokenizer.decode(inps), "t": tag})
                             else:
-                                this_sentence.append({"w": self.tokenizer.decode(inps), "t": tag, "l":lemma})
-                        else:
-                            this_sentence.append({"w":self.tokenizer.decode(inps).strip(), "t":tag, "l":lemma})                                                                   
-                    all_tagged_sentences.append({"lang":LANG_STR, "sent":  [ {"w":i["w"], "t":self.tags[LANG][i["t"]], "l":self.LemmaHandling.get_lemma_given_word_and_lemma_list_index(i["w"],i["l"])} for i in this_sentence] })
-                return all_tagged_sentences
-
-            # If the output is in TSV format
-            elif tag_config["output_tsv"]:
-                for input_ids, tags, lemmas in zip(all_input_ids, all_tags, all_lemmas):
-                    this_sentence=[]
-                    for inps, tag, lemma in zip(input_ids[1:], tags[1:], lemmas[1:]):
-                        if inps == self.tokenizer.sep_token_id or inps == self.tokenizer.pad_token_id:
-                            break
-                        if lemma == 0: # If there is no lemma here, that means we haven't reached the end of the word
-                            if len(this_sentence)>0:
-                                this_sentence[-1]["w"] += self.tokenizer.decode(inps)
-                            else:
-                                this_sentence.append({"w": self.tokenizer.decode(inps), "t": tag, "l":lemma})
-                        else:
-                            this_sentence.append({"w":self.tokenizer.decode(inps).strip(), "t":tag, "l":lemma})
-                    this_sentence=[ {"w":i["w"], "t":self.tags_str[LANG][i["t"]], "l":self.LemmaHandling.get_lemma_given_word_and_lemma_list_index(i["w"],i["l"])} for i in this_sentence]
-                    tag_config["write_output_to"].write(LANG_STR)
-                    for lin in this_sentence:
-                        tag_config["write_output_to"].write("\t")
-                        tag_config["write_output_to"].write(lin["w"])
-                        tag_config["write_output_to"].write("\t")
-                        tag_config["write_output_to"].write(lin["l"])
-                        tag_config["write_output_to"].write("\t")
-                        tag_config["write_output_to"].write(lin["t"])
+                                this_sentence.append({"w":self.tokenizer.decode(inps).strip(), "t":tag})
+                            
+                        json.dump({"lang":LANG_STR, "sent":[ {"w":i["w"], "t":self.tags[LANG][i["t"]]} for i in this_sentence]}, tag_config["write_output_to"])
                         tag_config["write_output_to"].write("\n")
-                    tag_config["write_output_to"].write("\n")
 
-            # If output format will be json
-            else:
-                for input_ids, tags, lemmas in zip(all_input_ids, all_tags, all_lemmas):
-                    this_sentence=[]
-                    for inps, tag, lemma in zip(input_ids[1:], tags[1:], lemmas[1:]):
-                        if inps == self.tokenizer.sep_token_id or inps == self.tokenizer.pad_token_id:
-                            break
-                        if lemma == 0: # If there is no lemma here, that means we haven't reached the end of the word
-                            if len(this_sentence)>0:
-                                this_sentence[-1]["w"] += self.tokenizer.decode(inps)
-                            else:
-                                this_sentence.append({"w": self.tokenizer.decode(inps), "t": tag, "l":lemma})
-                        else:
-                            this_sentence.append({"w":self.tokenizer.decode(inps).strip(), "t":tag, "l":lemma})
-                            
-                    json.dump({"lang":LANG_STR, "sent":[ {"w":i["w"], "t":self.tags[LANG][i["t"]], "l":self.LemmaHandling.get_lemma_given_word_and_lemma_list_index(i["w"],i["l"])} for i in this_sentence]}, tag_config["write_output_to"])
-                    tag_config["write_output_to"].write("\n")
 
     def _check_if_text_file_and_return_content(self, filepath):
         try:
@@ -575,7 +980,21 @@ class HumitTaggerModel(torch.nn.Module):
 
     @torch.no_grad()
     def tag(self, inp=None, **tag_config):
+
         self.eval()
+
+        if "lemmatise" in tag_config and tag_config["lemmatise"]==False:
+            tag_config["lemmatize"] = False
+            if "lemmatise" in tag_config:
+                del tag_config["lemmatise"]
+        else:
+            tag_config["lemmatize"] = True
+            if "lemmatise" in tag_config:
+                del tag_config["lemmatise"]
+
+        if "lemmatize" in tag_config and tag_config["lemmatize"]==False:
+            tag_config["lemmatize"] = False
+
         if "one_sentence_per_line" not in tag_config:
             tag_config["one_sentence_per_line"]=False
 
@@ -620,7 +1039,7 @@ class HumitTaggerModel(torch.nn.Module):
                         os.makedirs(os.path.dirname(out_path), exist_ok=True) 
                         if tag_config["one_sentence_per_line"]:
                             inp = [i for i in file_content.split("\n") if i!=""]
-                            inp = [i for i in inp if i!=""]
+                            inp = [" ".join(i.split()) for i in inp if i!=""]
                             with open(out_path, "w") as opened_file:
                                 tag_config["write_output_to"] = opened_file
                                 self.tag_sentence_list(inp, **tag_config)
@@ -631,8 +1050,8 @@ class HumitTaggerModel(torch.nn.Module):
                                 self.tag_sentence_list(inp, **tag_config)
                     else:
                         print (f"Could not properly open and read {input_path}.")
-
-            write_to.close()
+            if write_to is not sys.stdout and write_to is not sys.stderr:
+                write_to.close()
             return
 
         else:
@@ -650,7 +1069,7 @@ class HumitTaggerModel(torch.nn.Module):
             # Tag one sentence per line in a string
             if tag_config["one_sentence_per_line"]:
                 inp = [i for i in inp.split("\n") if i!=""]
-                inp = [self._preprocess_text(i) for i in inp if i!=""]
+                inp = [" ".join(self._preprocess_text(i).split()) for i in inp if i!=""]
                 return self.tag_sentence_list(inp, **tag_config)
 
             # identify sentences
@@ -660,7 +1079,7 @@ class HumitTaggerModel(torch.nn.Module):
         # Tag one sentence per list item
         elif type(inp) == list:
             inp=[i.strip() for i in inp]
-            inp=[self._preprocess_text(i) for i in inp if i!=""]
+            inp=[" ".join(self._preprocess_text(i).split()) for i in inp if i!=""]
             return self.tag_sentence_list(inp, **tag_config)
 
     def identify_language_sentence_list(self, lst, **tag_config):
@@ -703,9 +1122,12 @@ class HumitTaggerModel(torch.nn.Module):
 
     @torch.no_grad()
     def identify_language(self, inp=None, **tag_config):
+
         self.eval()
+
         if "one_sentence_per_line" not in tag_config:
             tag_config["one_sentence_per_line"]=False
+
         if "lang" in tag_config:
             del tag_config["lang"]
 
@@ -715,7 +1137,7 @@ class HumitTaggerModel(torch.nn.Module):
         if "lang_per_sentence" not in tag_config:
             tag_config["lang_per_sentence"] = False
 
-        elif tag_config["lang_per_sentence"]:
+        elif type(tag_config["lang_per_sentence"])==bool and tag_config["lang_per_sentence"]:
             tag_config["lang_per_sentence"] = True
 
         if "input_directory" in tag_config and "output_directory" in tag_config and "write_output_to" in tag_config and tag_config["write_output_to"]!=None:
@@ -771,7 +1193,7 @@ class HumitTaggerModel(torch.nn.Module):
                         torch.cuda.empty_cache()
 
                         if tag_config["write_output_to"]==None:
-                            general_output.extend([{"f":i[0], "l":self.config["id_to_lang"][i[1]]} for i in zip(file_names, langs)])
+                            general_output.extend([{"f":i[0], "lang":self.config["id_to_lang"][i[1]]} for i in zip(file_names, langs)])
                         elif tag_config["output_tsv"]:
                             for fil,lan in zip(file_names, langs):
                                 tag_config["write_output_to"].write(fil)
@@ -780,7 +1202,7 @@ class HumitTaggerModel(torch.nn.Module):
                                 tag_config["write_output_to"].write("\n")
                         else:
                             for fil,lan in zip(file_names, langs):
-                                json.dump({"f":fil, "l":self.config["id_to_lang"][lan]})
+                                json.dump({"f":fil, "lang":self.config["id_to_lang"][lan]})
                         file_names=[]
                         contents=[]
                     else:
@@ -801,7 +1223,7 @@ class HumitTaggerModel(torch.nn.Module):
                 torch.cuda.empty_cache()
                             
                 if tag_config["write_output_to"]==None:
-                    general_output.extend([{"f":i[0], "l":self.config["id_to_lang"][i[1]]} for i in zip(file_names, langs)])
+                    general_output.extend([{"f":i[0], "lang":self.config["id_to_lang"][i[1]]} for i in zip(file_names, langs)])
                 elif tag_config["output_tsv"]:                 
                     for fil,lan in zip(file_names, langs):       
                         tag_config["write_output_to"].write(fil)
@@ -810,7 +1232,7 @@ class HumitTaggerModel(torch.nn.Module):
                         tag_config["write_output_to"].write("\n")
                 else:   
                     for fil,lan in zip(file_names, langs):
-                        json.dump({"f":fil, "l":self.config["id_to_lang"][lan]})
+                        json.dump({"f":fil, "lang":self.config["id_to_lang"][lan]})
 
             return general_output if len(general_output)>0 else None
 
@@ -852,17 +1274,17 @@ class HumitTaggerModel(torch.nn.Module):
                                                 opened_file.write(lan)
                                                 opened_file.write("\n")
                                         else:
-                                            json.dump([{"s":sen, "l":lan} for sen,lan in zip(inp, out) ] , opened_file)
+                                            json.dump([{"s":sen, "lang":lan} for sen,lan in zip(inp, out) ] , opened_file)
                                     else:
                                         if tag_config["output_tsv"]:
                                             opened_file.write(out[0])
                                         else:
-                                            json.dump({"l":out[0]} , opened_file)
+                                            json.dump({"lang":out[0]} , opened_file)
                             else:
                                 if tag_config["lang_per_sentence"]:
-                                    general_output.extend([{"s":sen, "l":lan} for sen,lan in zip(inp, out) ])
+                                    general_output.extend([{"s":sen, "lang":lan} for sen,lan in zip(inp, out) ])
                                 else:
-                                    general_output.append({"f":input_path, "l":out[0]})
+                                    general_output.append({"f":input_path, "lang":out[0]})
 
                         # If there is an opened pipe already
                         else:
@@ -875,7 +1297,7 @@ class HumitTaggerModel(torch.nn.Module):
                                         tag_config["write_output_to"].write("\n")
                                     tag_config["write_output_to"].write("\n")
                                 else:
-                                    json.dump([{"s":sen, "l":lan} for sen,lan in zip(inp, out) ] , tag_config["write_output_to"])
+                                    json.dump([{"s":sen, "lang":lan} for sen,lan in zip(inp, out) ] , tag_config["write_output_to"])
                                     tag_config["write_output_to"].write("\n")
                             else:
                                 if tag_config["output_tsv"]:
@@ -884,7 +1306,7 @@ class HumitTaggerModel(torch.nn.Module):
                                     tag_config["write_output_to"].write(out[0])
                                     tag_config["write_output_to"].write("\n")
                                 else:    
-                                    json.dump({"f":input_path, "l":out[0]} , tag_config["write_output_to"])
+                                    json.dump({"f":input_path, "lang":out[0]} , tag_config["write_output_to"])
                                     tag_config["write_output_to"].write("\n")
 
                     else:
@@ -894,10 +1316,10 @@ class HumitTaggerModel(torch.nn.Module):
                             tag_config["write_output_to"].write("err")
                             tag_config["write_output_to"].write("\n")
                         else:    
-                            json.dump({"f":input_path, "l":"err"} , tag_config["write_output_to"])
+                            json.dump({"f":input_path, "lang":"err"} , tag_config["write_output_to"])
                             tag_config["write_output_to"].write("\n")
 
-            if tag_config["write_output_to"] and tag_config["write_output_to"]!=sys.stdout and tag_config["write_output_to"]!=sys.stderr:
+            if tag_config["write_output_to"] and tag_config["write_output_to"] is not sys.stdout and tag_config["write_output_to"] is not sys.stderr:
                 tag_config["write_output_to"].close() 
 
             return general_output if len(general_output)>0 else None
@@ -933,7 +1355,7 @@ class HumitTaggerModel(torch.nn.Module):
 
             # If return as list
             if tag_config["write_output_to"]==None:
-                return [{"s":i[0], "l": i[1]} for i in zip(inp, out)]
+                return [{"s":i[0], "lang": i[1]} for i in zip(inp, out)]
 
             if tag_config["output_tsv"]:
                 for sen,lan in zip(inp, out):
@@ -942,7 +1364,7 @@ class HumitTaggerModel(torch.nn.Module):
                     tag_config["write_output_to"].write(out)
                     tag_config["write_output_to"].write("\n")
             else:
-                json.dump([{"s":sen, "l":lan} for sen,lan in zip(inp, out) ] , tag_config["write_output_to"])
+                json.dump([{"s":sen, "lang":lan} for sen,lan in zip(inp, out) ] , tag_config["write_output_to"])
 
             return
 
@@ -954,7 +1376,7 @@ class HumitTaggerModel(torch.nn.Module):
 
             # If return as list 
             if tag_config["write_output_to"]==None:
-                return [{"s":i[0], "l": i[1]} for i in zip(inp, out)]
+                return [{"s":i[0], "lang": i[1]} for i in zip(inp, out)]
 
             if tag_config["output_tsv"]:    
                 for sen,lan in zip(inp, out):    
@@ -963,7 +1385,7 @@ class HumitTaggerModel(torch.nn.Module):
                     tag_config["write_output_to"].write(lan)
                     tag_config["write_output_to"].write("\n")
             else:    
-                json.dump([{"s":sen, "l":lan} for sen,lan in zip(inp, out) ] , tag_config["write_output_to"])
+                json.dump([{"s":sen, "lang":lan} for sen,lan in zip(inp, out) ] , tag_config["write_output_to"])
 
             return