Update app.py
Browse files
app.py
CHANGED
@@ -5,13 +5,11 @@ import torch
|
|
5 |
|
6 |
# Load the model and config when the script starts
|
7 |
config = PeftConfig.from_pretrained("phearion/bigbrain-v0.0.1")
|
|
|
8 |
model = AutoModelForCausalLM.from_pretrained("mistralai/Mistral-7B-v0.1")
|
9 |
model = PeftModel.from_pretrained(model,
|
10 |
"phearion/bigbrain-v0.0.1")
|
11 |
|
12 |
-
# Convert the model to TorchScript
|
13 |
-
scripted_model = torch.jit.script(model)
|
14 |
-
|
15 |
# Load the tokenizer
|
16 |
tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-v0.1")
|
17 |
|
@@ -22,7 +20,7 @@ def greet(text):
|
|
22 |
|
23 |
# Use torch.no_grad to disable gradient calculation
|
24 |
with torch.no_grad():
|
25 |
-
output_tokens =
|
26 |
|
27 |
return tokenizer.decode(output_tokens[0], skip_special_tokens=True)
|
28 |
|
|
|
5 |
|
6 |
# Load the model and config when the script starts
|
7 |
config = PeftConfig.from_pretrained("phearion/bigbrain-v0.0.1")
|
8 |
+
|
9 |
model = AutoModelForCausalLM.from_pretrained("mistralai/Mistral-7B-v0.1")
|
10 |
model = PeftModel.from_pretrained(model,
|
11 |
"phearion/bigbrain-v0.0.1")
|
12 |
|
|
|
|
|
|
|
13 |
# Load the tokenizer
|
14 |
tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-v0.1")
|
15 |
|
|
|
20 |
|
21 |
# Use torch.no_grad to disable gradient calculation
|
22 |
with torch.no_grad():
|
23 |
+
output_tokens = model.generate(**batch, max_new_tokens=15)
|
24 |
|
25 |
return tokenizer.decode(output_tokens[0], skip_special_tokens=True)
|
26 |
|