Jeff Myers II commited on
Commit
405db4b
·
1 Parent(s): cb3e313

Removed task arg from pipeline

Browse files
Files changed (1) hide show
  1. Gemma.py +1 -1
Gemma.py CHANGED
@@ -21,7 +21,7 @@ class GemmaLLM:
21
  # model = AutoModelForCausalLM.from_pretrained(model_id, quantization_config=quant_config)
22
  # tokenizer = AutoTokenizer.from_pretrained(model_id)
23
 
24
- self.model = pipeline("text-generation", model_id)
25
 
26
  @spaces.GPU
27
  def generate(self, message) -> str:
 
21
  # model = AutoModelForCausalLM.from_pretrained(model_id, quantization_config=quant_config)
22
  # tokenizer = AutoTokenizer.from_pretrained(model_id)
23
 
24
+ self.model = pipeline(model=model_id)
25
 
26
  @spaces.GPU
27
  def generate(self, message) -> str: