Update app.py
Browse files
app.py
CHANGED
@@ -3,7 +3,6 @@ import gradio as gr
|
|
3 |
from transformers import LlamaTokenizer, AutoModelForCausalLM
|
4 |
import torch
|
5 |
import json
|
6 |
-
from datetime import datetime
|
7 |
|
8 |
# Set Hugging Face Token for Authentication (ensure it's set in your environment)
|
9 |
HUGGINGFACE_TOKEN = os.getenv("HUGGINGFACE_TOKEN")
|
@@ -14,7 +13,7 @@ tokenizer = LlamaTokenizer.from_pretrained(MODEL_NAME, token=HUGGINGFACE_TOKEN)
|
|
14 |
model = AutoModelForCausalLM.from_pretrained(
|
15 |
MODEL_NAME,
|
16 |
token=HUGGINGFACE_TOKEN,
|
17 |
-
device_map="cpu"
|
18 |
)
|
19 |
|
20 |
# Load Llama Guard for content moderation on CPU
|
|
|
3 |
from transformers import LlamaTokenizer, AutoModelForCausalLM
|
4 |
import torch
|
5 |
import json
|
|
|
6 |
|
7 |
# Set Hugging Face Token for Authentication (ensure it's set in your environment)
|
8 |
HUGGINGFACE_TOKEN = os.getenv("HUGGINGFACE_TOKEN")
|
|
|
13 |
model = AutoModelForCausalLM.from_pretrained(
|
14 |
MODEL_NAME,
|
15 |
token=HUGGINGFACE_TOKEN,
|
16 |
+
device_map="cpu"
|
17 |
)
|
18 |
|
19 |
# Load Llama Guard for content moderation on CPU
|