Norod78 commited on
Commit
89f84ba
1 Parent(s): c7dd9b3

Set a repetition_penalty constant as 1.8

Browse files

Following the discussion here: https://huggingface.co/spaces/Tonic/Aya/discussions/2

Files changed (1) hide show
  1. app.py +6 -1
app.py CHANGED
@@ -14,11 +14,16 @@ checkpoint = "CohereForAI/aya-101"
14
  tokenizer = AutoTokenizer.from_pretrained(checkpoint)
15
  model = AutoModelForSeq2SeqLM.from_pretrained(checkpoint, torch_dtype=torch.bfloat16, low_cpu_mem_usage=True, device_map=device)
16
 
 
 
 
 
 
17
  @spaces.GPU
18
  def aya(text, max_new_tokens):
19
  model.to(device)
20
  inputs = tokenizer.encode(text, return_tensors="pt").to(device)
21
- outputs = model.generate(inputs, max_new_tokens=max_new_tokens)
22
  translation = tokenizer.decode(outputs[0], skip_special_tokens=True)
23
  return translation
24
 
 
14
  tokenizer = AutoTokenizer.from_pretrained(checkpoint)
15
  model = AutoModelForSeq2SeqLM.from_pretrained(checkpoint, torch_dtype=torch.bfloat16, low_cpu_mem_usage=True, device_map=device)
16
 
17
+ #Set a the value of the repetition penalty
18
+ #The higher the value, the less repetitive the generated text will be
19
+ #Note that `repetition_penalty` has to be a strictly positive float
20
+ repetition_penalty = 1.8
21
+
22
  @spaces.GPU
23
  def aya(text, max_new_tokens):
24
  model.to(device)
25
  inputs = tokenizer.encode(text, return_tensors="pt").to(device)
26
+ outputs = model.generate(inputs, max_new_tokens=max_new_tokens, repetition_penalty=repetition_penalty)
27
  translation = tokenizer.decode(outputs[0], skip_special_tokens=True)
28
  return translation
29