Update app.py (#1)

- Update app.py (472cb7ecbc635f2f109f555f6267af14cbc4dfc9)


Co-authored-by: ilyas aqit <[email protected]>

app.py

@@ -31,25 +31,29 @@ model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME).to(device)
 tokenizer = NllbTokenizer.from_pretrained(MODEL_NAME)
 
 
-def translate(text, source_lang, target_lang, max_length=238, num_beams=4, repetition_penalty=1.0):
+def translate(text, source_lang="English", target_lang="Tachelhit/Central Atlas Tamazight",
+              max_length=238, num_beams=4, repetition_penalty=1.0):
     """
-    Translate text from source language to target language
+    Translate multi-line text while preserving line breaks.
+    Each line is translated independently.
     """
-    print(text)
-
-    tokenizer.src_lang = NLLB_LANG_MAPPING[source_lang]
-
-    inputs = tokenizer(text, return_tensors="pt").to(model.device)
-    translated_tokens = model.generate(
-        **inputs,
-        forced_bos_token_id=tokenizer.convert_tokens_to_ids(NLLB_LANG_MAPPING[target_lang]),
-        max_length=max_length,
-        num_beams=num_beams,
-        repetition_penalty=float(repetition_penalty),
-    )
-    translation = tokenizer.batch_decode(translated_tokens, skip_special_tokens=True)[0]
-
-    return translation
+    translations = []
+    for line in text.split("\n"):
+        if line.strip() == "":
+            translations.append("")  # preserve empty lines
+        else:
+            tokenizer.src_lang = NLLB_LANG_MAPPING[source_lang]
+            inputs = tokenizer(line, return_tensors="pt").to(model.device)
+            translated_tokens = model.generate(
+                **inputs,
+                forced_bos_token_id=tokenizer.convert_tokens_to_ids(NLLB_LANG_MAPPING[target_lang]),
+                max_length=max_length,
+                num_beams=num_beams,
+                repetition_penalty=float(repetition_penalty),
+            )
+            translation = tokenizer.batch_decode(translated_tokens, skip_special_tokens=True)[0]
+            translations.append(translation)
+    return "\n".join(translations)
 
 
 gradio_ui= gr.Interface(