TiberiuCristianLeon commited on
Commit
fb22ed5
·
verified ·
1 Parent(s): e8cfbc3

Update src/translate/Translate.py

Browse files
Files changed (1) hide show
  1. src/translate/Translate.py +7 -3
src/translate/Translate.py CHANGED
@@ -53,17 +53,21 @@ def paraphraseTranslateMethod(requestValue: str, model: str):
53
  return " ".join(result_value).strip(), model
54
 
55
  def gemma(requestValue: str, model: str = 'Gargaz/gemma-2b-romanian-better'):
 
56
  prompt = f"Translate this to Romanian using a formal tone. Only return the translation: {requestValue}"
57
  messages = [{"role": "user", "content": f"Translate this text to Romanian using a formal tone. Only return the translated text: {requestValue}"}]
58
  if '/' not in model:
59
  model = 'Gargaz/gemma-2b-romanian-better'
 
60
  pipe = pipeline(
61
  "text-generation",
62
  model=model,
63
  device=-1,
64
- max_new_tokens=256, # Keep short to reduce verbosity
65
  do_sample=False # Use greedy decoding for determinism
66
  )
67
  output = pipe(messages, num_return_sequences=1, return_full_text=False)
68
- # return output[0]["generated_text"].strip(), model
69
- return output, model
 
 
 
53
  return " ".join(result_value).strip(), model
54
 
55
  def gemma(requestValue: str, model: str = 'Gargaz/gemma-2b-romanian-better'):
56
+ requestValue = requestValue.replace('\n', ' ')
57
  prompt = f"Translate this to Romanian using a formal tone. Only return the translation: {requestValue}"
58
  messages = [{"role": "user", "content": f"Translate this text to Romanian using a formal tone. Only return the translated text: {requestValue}"}]
59
  if '/' not in model:
60
  model = 'Gargaz/gemma-2b-romanian-better'
61
+ max_new_tokens = len(requestValue) + len(requestValue) * 0.2
62
  pipe = pipeline(
63
  "text-generation",
64
  model=model,
65
  device=-1,
66
+ max_new_tokens=max_new_tokens, # Keep short to reduce verbosity
67
  do_sample=False # Use greedy decoding for determinism
68
  )
69
  output = pipe(messages, num_return_sequences=1, return_full_text=False)
70
+ generated_text = output[0]["generated_text"]
71
+ result = generated_text.split('\n', 1)[0].strip()
72
+ return result, model
73
+ # return output, model