Steph254 commited on
Commit
5e6f6cd
·
verified ·
1 Parent(s): 6f7715b

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +1 -2
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" # Force CPU usage
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