SurajSingh commited on
Commit
c58e248
·
verified ·
1 Parent(s): 44db241

Update main.py

Browse files
Files changed (1) hide show
  1. main.py +1 -4
main.py CHANGED
@@ -1,7 +1,4 @@
1
  from fastapi import FastAPI
2
- import os
3
- os.environ["TRANSFORMERS_CACHE"] = "./checkpoint"
4
-
5
 
6
  from transformers import LineByLineTextDataset
7
  from transformers import DataCollatorForLanguageModeling
@@ -28,7 +25,7 @@ def load_tokenizer(checkpoint):
28
 
29
  model_path = r'./checkpoint/'
30
  model = load_model(model_path)
31
- tokenizer = load_tokenizer('gpt2')
32
 
33
 
34
  def generate_text(sequence, max_new_tokens):
 
1
  from fastapi import FastAPI
 
 
 
2
 
3
  from transformers import LineByLineTextDataset
4
  from transformers import DataCollatorForLanguageModeling
 
25
 
26
  model_path = r'./checkpoint/'
27
  model = load_model(model_path)
28
+ tokenizer = load_tokenizer('./tokenizer/')
29
 
30
 
31
  def generate_text(sequence, max_new_tokens):