TiberiuCristianLeon commited on
Commit
e7d7478
·
verified ·
1 Parent(s): 6b871de

Update src/translate/Translate.py

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