Update src/translate/Translate.py

src/translate/Translate.py

@@ -74,12 +74,12 @@ 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
-    model = model if '/' in model else 'Gargaz/gemma-2b-romanian-better'
+    model_name = model if '/' in model else 'Gargaz/gemma-2b-romanian-better'
     # limit max_new_tokens to 150% of the requestValue  
     prompt = f"Translate this text to Romanian: {requestValue}"
     
-    tokenizer = AutoTokenizer.from_pretrained("Gargaz/gemma-2b-romanian-better")
-    model = AutoModelForCausalLM.from_pretrained("Gargaz/gemma-2b-romanian-better").to(device)
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    model = AutoModelForCausalLM.from_pretrained(model_name).to(device)
     input_ids = tokenizer.encode(requestValue, add_special_tokens=True)
     num_tokens = len(input_ids)
     # Estimate output length (e.g., 50% longer)
@@ -88,14 +88,18 @@ def gemma_direct(requestValue: str, model: str = 'Gargaz/gemma-2b-romanian-bette
        
     messages = [{"role": "user", "content": prompt}]
     
-    inputs = tokenizer.apply_chat_template(
-        messages,
-        add_generation_prompt=True,
-        tokenize=True,
-        return_dict=True,
-        return_tensors="pt",
-    ).to(device)
-
-    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
+    try:
+        inputs = tokenizer.apply_chat_template(
+            messages,
+            add_generation_prompt=True,
+            tokenize=True,
+            return_dict=True,
+            return_tensors="pt",
+        ).to(device)
+    
+        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
+    except Exception as error:
+        return error
+