mrm8488 commited on
Commit
5154ed3
1 Parent(s): 55036a1

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +2 -2
app.py CHANGED
@@ -11,12 +11,12 @@ CKPT = "mrm8488/santacoder-finetuned-the-stack-swift"
11
 
12
  examples = [[EXAMPLE_0, 9, 0.6, 42], [EXAMPLE_1, 114, 0.6, 42]]
13
  tokenizer = AutoTokenizer.from_pretrained(CKPT)
14
- model = AutoModelForCausalLM.from_pretrained(CKPT, trust_remote_code=True).to("cuda")
15
 
16
 
17
  def code_generation(gen_prompt, max_tokens, temperature=0.6, seed=42):
18
  set_seed(seed)
19
- pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, device=0)
20
  generated_text = pipe(gen_prompt, do_sample=True, top_p=0.95, temperature=temperature, max_new_tokens=max_tokens)[0]['generated_text']
21
  return generated_text
22
 
 
11
 
12
  examples = [[EXAMPLE_0, 9, 0.6, 42], [EXAMPLE_1, 114, 0.6, 42]]
13
  tokenizer = AutoTokenizer.from_pretrained(CKPT)
14
+ model = AutoModelForCausalLM.from_pretrained(CKPT, trust_remote_code=True)
15
 
16
 
17
  def code_generation(gen_prompt, max_tokens, temperature=0.6, seed=42):
18
  set_seed(seed)
19
+ pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
20
  generated_text = pipe(gen_prompt, do_sample=True, top_p=0.95, temperature=temperature, max_new_tokens=max_tokens)[0]['generated_text']
21
  return generated_text
22