Skip to content
Snippets Groups Projects
Commit 486bff9d authored by Ludwig Forsberg's avatar Ludwig Forsberg
Browse files

setup remote fetch to model

parent 320dc33c
No related branches found
No related tags found
No related merge requests found
...@@ -6,7 +6,7 @@ from transformers.models.bart.modeling_bart import shift_tokens_right ...@@ -6,7 +6,7 @@ from transformers.models.bart.modeling_bart import shift_tokens_right
# Save the model and tokenizer as local variables to use them in the predict_query function # Save the model and tokenizer as local variables to use them in the predict_query function
def setup(): def setup():
# Import model from ./trained-models/blackbox # Import model from ./trained-models/blackbox
model = BartForConditionalGeneration.from_pretrained("./trained-models/blackbox") model = BartForConditionalGeneration.from_pretrained("ludfo774/sparql-bart-append-1")
tokenizer = BartTokenizer.from_pretrained("facebook/bart-base") tokenizer = BartTokenizer.from_pretrained("facebook/bart-base")
model.to('cuda') model.to('cuda')
return model, tokenizer return model, tokenizer
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Please register or to comment