PhantHive commited on
Commit
98e2a1f
·
verified ·
1 Parent(s): ef29b9f

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +3 -3
app.py CHANGED
@@ -12,8 +12,7 @@ else:
12
  print("GPU is not available, using CPU.")
13
 
14
  # Load the model and config when the script starts
15
- peft_model_id =
16
- "phearion/bigbrain-v0.0.1"
17
  config = PeftConfig.from_pretrained(peft_model_id)
18
  model = AutoModelForCausalLM.from_pretrained(config.base_model_name_or_path, return_dict=True)
19
  tokenizer = AutoTokenizer.from_pretrained(config.base_model_name_or_path)
@@ -26,7 +25,8 @@ def greet(text):
26
  batch = tokenizer(f"\"{text}\" ->: ", return_tensors='pt')
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
 
 
12
  print("GPU is not available, using CPU.")
13
 
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, return_dict=True)
18
  tokenizer = AutoTokenizer.from_pretrained(config.base_model_name_or_path)
 
25
  batch = tokenizer(f"\"{text}\" ->: ", return_tensors='pt')
26
 
27
  # Use torch.no_grad to disable gradient calculation
28
+ with torch
29
+ .no_grad():
30
  output_tokens = model.generate(**batch, do_sample=True, max_new_tokens=15
31
  )
32