selbl commited on
Commit
1cbf66f
1 Parent(s): 6ad070d

Update StreamlitModel.py

Browse files
Files changed (1) hide show
  1. StreamlitModel.py +1 -1
StreamlitModel.py CHANGED
@@ -104,7 +104,7 @@ def TextGeneration(prompt,prof=False,parts=True):
104
  configuration = GPT2Config()
105
  gpt_model = GPT2_Model(configuration).to(device)
106
  #gpt_model.load_state_dict(torch.load('GPT-Trained-Model-Prod.pt'))
107
- state_dict = torch.hub.load_state_dict_from_url(r'https://github.com/Selbl/LyricGeneration/raw/main/GPT-Trained-Model-Prod.pt?download=', weights_only=True)
108
  gpt_model.load_state_dict(state_dict)
109
  #Load tokenizer
110
  tokenizer = GPT2Tokenizer.from_pretrained('gpt2', pad_token='<|pad|>')
 
104
  configuration = GPT2Config()
105
  gpt_model = GPT2_Model(configuration).to(device)
106
  #gpt_model.load_state_dict(torch.load('GPT-Trained-Model-Prod.pt'))
107
+ state_dict = torch.hub.load_state_dict_from_url(r'https://github.com/Selbl/LyricGeneration/raw/main/GPT-Trained-Model-Prod.pt?download=')
108
  gpt_model.load_state_dict(state_dict)
109
  #Load tokenizer
110
  tokenizer = GPT2Tokenizer.from_pretrained('gpt2', pad_token='<|pad|>')