kz209 commited on
Commit
da4acea
·
1 Parent(s): fcae4c8
Files changed (1) hide show
  1. utils/model.py +2 -2
utils/model.py CHANGED
@@ -4,7 +4,7 @@ import torch
4
 
5
 
6
  class Model():
7
- def __init__(self, model="tiiuae/falcon-7b-instruct") -> None:
8
  pass
9
  self.tokenizer = AutoTokenizer.from_pretrained(model)
10
  self.pipeline = transformers.pipeline(
@@ -20,7 +20,7 @@ class Model():
20
  sequences = self.pipeline(
21
  content,
22
  max_new_tokens=max_length,
23
- do_sample=False,
24
  temperature=temp,
25
  num_return_sequences=1,
26
  eos_token_id=self.tokenizer.eos_token_id,
 
4
 
5
 
6
  class Model():
7
+ def __init__(self, model="mistralai/Mistral-7B-Instruct-v0.1") -> None:
8
  pass
9
  self.tokenizer = AutoTokenizer.from_pretrained(model)
10
  self.pipeline = transformers.pipeline(
 
20
  sequences = self.pipeline(
21
  content,
22
  max_new_tokens=max_length,
23
+ do_sample=True,
24
  temperature=temp,
25
  num_return_sequences=1,
26
  eos_token_id=self.tokenizer.eos_token_id,