Update app/model/model.py

app/model/model.py

@@ -28,9 +28,11 @@ class LLM:
         tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=True)
         if tokenizer.pad_token_id is None:
             tokenizer.pad_token_id = tokenizer.eos_token_id
+        print("Model and tokenizer loaded.")
         return model, tokenizer
 
     def language_detection(self, input_text):
+        print(f"### Input text\n{input_text}")
         # Prompt with one shot for each language
         prompt = f"""Identify the language of the following sentences. Options: 'english', 'español', 'française' .
             * <Identity theft is not a joke, millions of families suffer every year>(english)
@@ -39,9 +41,11 @@ class LLM:
             * <{input_text}>"""
         # Generation and extraction of the language tag
         answer_ids = self.model.generate(**self.tokenizer([prompt], return_tensors="pt"), max_new_tokens=10)
-        answer = self.tokenizer.batch_decode(answer_ids, skip_special_tokens=False)[0].split(prompt)[1]
+        answer = self.tokenizer.batch_decode(answer_ids, skip_special_tokens=False)[0]
+        print(answer)
+        generation = answer.split(prompt)[1]
         pattern = r'\b(?:' + '|'.join(map(re.escape, self.lang_codes.keys())) + r')\b'
-        lang = re.search(pattern, answer, flags=re.IGNORECASE)
+        lang = re.search(pattern, generation, flags=re.IGNORECASE)
         # Returns tag identified or 'unk' if none is detected
         return self.lang_codes[lang.group()] if lang else self.lang_codes["unknown"]