Update app.py
Browse files
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
|
|
|
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 |
|