PhantHive commited on
Commit
2b702db
·
verified ·
1 Parent(s): 187b236

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +2 -4
app.py CHANGED
@@ -14,9 +14,7 @@ else:
14
  # Load the model and config when the script starts
15
  peft_model_id = "phearion/bigbrain-v0.0.1"
16
  config = PeftConfig.from_pretrained(peft_model_id)
17
- model = AutoModelForCausalLM.from_pretrained(
18
- config.base_model_name_or_path,
19
- torch_dtype=torch.bfloat16)
20
  tokenizer = AutoTokenizer.from_pretrained(config.base_model_name_or_path)
21
 
22
  # Load the Lora model
@@ -27,7 +25,7 @@ def greet(text):
27
 
28
  # Use torch.no_grad to disable gradient calculation
29
  with torch.no_grad():
30
- output_tokens = model.generate(**batch, do_sample=True, max_new_tokens=15)
31
 
32
  return tokenizer.decode(output_tokens[0], skip_special_tokens=True)
33
 
 
14
  # Load the model and config when the script starts
15
  peft_model_id = "phearion/bigbrain-v0.0.1"
16
  config = PeftConfig.from_pretrained(peft_model_id)
17
+ model = AutoModelForCausalLM.from_pretrained(config.base_model_name_or_path)
 
 
18
  tokenizer = AutoTokenizer.from_pretrained(config.base_model_name_or_path)
19
 
20
  # Load the Lora model
 
25
 
26
  # Use torch.no_grad to disable gradient calculation
27
  with torch.no_grad():
28
+ output_tokens = model.generate(**batch, do_sample=True, max_new_tokens=20)
29
 
30
  return tokenizer.decode(output_tokens[0], skip_special_tokens=True)
31