Update src/translate/Translate.py

src/translate/Translate.py

@@ -60,17 +60,20 @@ def gemma(requestValue: str, model: str = 'Gargaz/gemma-2b-romanian-better'):
         model = 'Gargaz/gemma-2b-romanian-better'
     # limit max_new_tokens to 150% of the requestValue  
     max_new_tokens = int(len(requestValue) + len(requestValue) * 0.5)
-    pipe = pipeline(
+    try:
+        pipe = pipeline(
         "text-generation",
         model=model,
         device=-1,
         max_new_tokens=max_new_tokens,   # Keep short to reduce verbosity
         do_sample=False  # Use greedy decoding for determinism
-    ) 
-    output = pipe(messages, num_return_sequences=1, return_full_text=False)
-    generated_text = output[0]["generated_text"]
-    result = generated_text.split('\n', 1)[0].strip()
-    return result, model
+        )
+        output = pipe(messages, num_return_sequences=1, return_full_text=False)
+        generated_text = output[0]["generated_text"]
+        result = generated_text.split('\n', 1)[0] if '\n' in generated_text else generated_text
+        return result.strip()
+    except Exception as error:
+        return error
 
 def gemma_direct(requestValue: str, model: str = 'Gargaz/gemma-2b-romanian-better'):
     # Load model directly
@@ -99,7 +102,8 @@ def gemma_direct(requestValue: str, model: str = 'Gargaz/gemma-2b-romanian-bette
     
         outputs = model.generate(**inputs, max_new_tokens=max_new_tokens)
         response = tokenizer.decode(outputs[0][inputs["input_ids"].shape[-1]:], skip_special_tokens=True)
-        return response.strip()
+        result = response.split('\n', 1)[0] if '\n' in response else response 
+        return result.strip()
     except Exception as error:
         return error