Tonic commited on
Commit
aa9a886
1 Parent(s): 69e7b01

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +3 -0
app.py CHANGED
@@ -33,9 +33,12 @@ class TuluChatBot:
33
  prompt = self.format_prompt(user_message)
34
  inputs = self.tokenizer(prompt, return_tensors='pt', add_special_tokens=False)
35
  input_ids = inputs["input_ids"].to(self.model.device)
 
 
36
 
37
  output_ids = self.model.generate(
38
  input_ids,
 
39
  max_length=input_ids.shape[1] + max_new_tokens,
40
  temperature=temperature,
41
  top_p=top_p,
 
33
  prompt = self.format_prompt(user_message)
34
  inputs = self.tokenizer(prompt, return_tensors='pt', add_special_tokens=False)
35
  input_ids = inputs["input_ids"].to(self.model.device)
36
+ attention_mask = inputs["attention_mask"].to(self.model.device)
37
+
38
 
39
  output_ids = self.model.generate(
40
  input_ids,
41
+ attention_mask=attention_mask,
42
  max_length=input_ids.shape[1] + max_new_tokens,
43
  temperature=temperature,
44
  top_p=top_p,