Spaces:

Luigi
/

Input-Method-Acceleration

Running on Zero

Luigi commited on 16 days ago

Commit

5799add

1 Parent(s): abaf626

enable trust_remote_code=True

Files changed (1) hide show

app.py CHANGED Viewed

@@ -19,7 +19,7 @@ MODEL_LIST = [
 @lru_cache(maxsize=None)
 def get_pipeline(model_name):
     tok = AutoTokenizer.from_pretrained(model_name)
-    mdl = AutoModelForCausalLM.from_pretrained(model_name, weights_only=False)
     mdl.to("cuda")
     return pipeline("text-generation", model=mdl, tokenizer=tok, device=0)

 @lru_cache(maxsize=None)
 def get_pipeline(model_name):
     tok = AutoTokenizer.from_pretrained(model_name)
+    mdl = AutoModelForCausalLM.from_pretrained(model_name, weights_only=False, trust_remote_code=True)
     mdl.to("cuda")
     return pipeline("text-generation", model=mdl, tokenizer=tok, device=0)