Update code/inference.py
Browse files- code/inference.py +1 -0
code/inference.py
CHANGED
@@ -3,6 +3,7 @@ import torch
|
|
3 |
from transformers import AutoModelForCausalLM, AutoTokenizer
|
4 |
from typing import List, Dict
|
5 |
from accelerate import load_checkpoint_and_dispatch
|
|
|
6 |
# Global variables to persist the model and tokenizer between invocations
|
7 |
model = None
|
8 |
tokenizer = None
|
|
|
3 |
from transformers import AutoModelForCausalLM, AutoTokenizer
|
4 |
from typing import List, Dict
|
5 |
from accelerate import load_checkpoint_and_dispatch
|
6 |
+
import os
|
7 |
# Global variables to persist the model and tokenizer between invocations
|
8 |
model = None
|
9 |
tokenizer = None
|