Ben Burtenshaw commited on
Commit
956b1a1
·
1 Parent(s): e838258

use transformers

Browse files
Files changed (1) hide show
  1. prometheus_pipeline.py +2 -2
prometheus_pipeline.py CHANGED
@@ -7,7 +7,7 @@ import time
7
  from distilabel.pipeline import Pipeline
8
  from distilabel.steps import KeepColumns, LoadHubDataset
9
  from distilabel.steps.tasks import PrometheusEval
10
- from distilabel.llms import vLLM
11
 
12
  if __name__ == "__main__":
13
  start_time = time.time()
@@ -22,7 +22,7 @@ if __name__ == "__main__":
22
 
23
  task = PrometheusEval(
24
  name="task",
25
- llm=vLLM(
26
  model="prometheus-eval/prometheus-7b-v2.0",
27
  chat_template="[INST] {{ messages[0]['content'] }}\n{{ messages[1]['content'] }}[/INST]",
28
  ),
 
7
  from distilabel.pipeline import Pipeline
8
  from distilabel.steps import KeepColumns, LoadHubDataset
9
  from distilabel.steps.tasks import PrometheusEval
10
+ from distilabel.llms import TransformersLLM
11
 
12
  if __name__ == "__main__":
13
  start_time = time.time()
 
22
 
23
  task = PrometheusEval(
24
  name="task",
25
+ llm=TransformersLLM(
26
  model="prometheus-eval/prometheus-7b-v2.0",
27
  chat_template="[INST] {{ messages[0]['content'] }}\n{{ messages[1]['content'] }}[/INST]",
28
  ),