Handle new lines

app.py

@@ -21,23 +21,30 @@ model = MistralForCausalLM.from_pretrained("mistralai/Mistral-7B-v0.1",
 model = PeftModel.from_pretrained(model, "lang-uk/dragoman").to("cuda")
 tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-v0.1", use_fast=False, add_bos_token=False)
 
-@spaces.GPU(duration=20)
+@spaces.GPU(duration=30)
 def translate(input_text):
-    # iteratively generate
-    input_text = input_text.strip()
-    input_text = f"[INST] {input_text} [/INST]"
-    inputs = tokenizer([input_text], return_tensors="pt").to(model.device)
-
-    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
-    generation_kwargs = dict(inputs, streamer=streamer, max_new_tokens=200)
-
-    thread = Thread(target=model.generate, kwargs=generation_kwargs)
-
-    thread.start()
-    
     generated_text = ""
-    for new_text in streamer:
-        generated_text += new_text
+    input_text = input_text.strip()
+    for chunk in input_text.split("\n"):
+        if not chunk:
+            generated_text += "\n"
+            yield generated_text
+            continue
+        chunk = f"[INST] {chunk} [/INST]"
+        inputs = tokenizer([chunk], return_tensors="pt").to(model.device)
+
+        streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
+        generation_kwargs = dict(inputs, streamer=streamer, max_new_tokens=200)
+
+        thread = Thread(target=model.generate, kwargs=generation_kwargs)
+
+        thread.start()
+        
+        for new_text in streamer:
+            generated_text += new_text
+            yield generated_text
+        
+        generated_text += "\n"
         yield generated_text