Update app.py
Browse files
app.py
CHANGED
@@ -6,9 +6,9 @@ import torch
|
|
6 |
# Initialize FastAPI app
|
7 |
app = FastAPI()
|
8 |
|
9 |
-
# Load pre-trained DistilGPT-2 model and tokenizer
|
10 |
model_name = "distilgpt2" # Smaller GPT-2 model
|
11 |
-
model = AutoModelForCausalLM.from_pretrained(model_name)
|
12 |
tokenizer = AutoTokenizer.from_pretrained(model_name)
|
13 |
|
14 |
# Pydantic model for request body
|
|
|
6 |
# Initialize FastAPI app
|
7 |
app = FastAPI()
|
8 |
|
9 |
+
# Load pre-trained DistilGPT-2 model and tokenizer, using from_tf=True to load TensorFlow weights
|
10 |
model_name = "distilgpt2" # Smaller GPT-2 model
|
11 |
+
model = AutoModelForCausalLM.from_pretrained(model_name, from_tf=True) # Use from_tf=True
|
12 |
tokenizer = AutoTokenizer.from_pretrained(model_name)
|
13 |
|
14 |
# Pydantic model for request body
|