Update src/translate/Translate.py

src/translate/Translate.py

@@ -74,14 +74,23 @@ def gemma(requestValue: str, model: str = 'Gargaz/gemma-2b-romanian-better'):
 
 def gemma_direct(requestValue: str, model: str = 'Gargaz/gemma-2b-romanian-better'):
     # Load model directly
-    if '/' not in model:
-        model = 'Gargaz/gemma-2b-romanian-better'
+    model = model if '/' in model else 'Gargaz/gemma-2b-romanian-better'
     # limit max_new_tokens to 150% of the requestValue  
-    max_new_tokens = int(len(requestValue) + len(requestValue) * 0.5)
-    max_new_tokens = max_new_tokens if max_new_tokens % 2 == 0 else max_new_tokens + 1
-    messages = [{"role": "user", "content": f"Translate this text to Romanian: {requestValue}"}]
+    prompt = f"Translate this text to Romanian: {request_value}"
+    
+    input_ids = tokenizer.encode(request_value, add_special_tokens=True)
+    num_tokens = len(input_ids)
+    # Estimate output length (e.g., 50% longer)
+    max_new_tokens = int(num_tokens * 1.5)
+    max_new_tokens += max_new_tokens % 2  # ensure it's even
+    
+    # Token count estimation and safety check
+    # max_new_tokens = int(len(request_value) * 1.5)
+    # max_new_tokens += max_new_tokens % 2  # ensure it's even
+    
+    messages = [{"role": "user", "content": prompt]
     tokenizer = AutoTokenizer.from_pretrained("Gargaz/gemma-2b-romanian-better")
-    model = AutoModelForCausalLM.from_pretrained("Gargaz/gemma-2b-romanian-better")
+    model = AutoModelForCausalLM.from_pretrained("Gargaz/gemma-2b-romanian-better").to(device)
     
     inputs = tokenizer.apply_chat_template(
         messages,
@@ -91,6 +100,6 @@ def gemma_direct(requestValue: str, model: str = 'Gargaz/gemma-2b-romanian-bette
         return_tensors="pt",
     ).to(device)
 
-    outputs = model.generate(**inputs, max_new_tokens=max_tokens)
+    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, model