Update app.py

app.py

@@ -1,6 +1,6 @@
 import os
 import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
+from transformers import AutoTokenizer, AutoModelForCausalLM
 import gradio as gr
 import spaces
 
@@ -8,12 +8,10 @@ huggingface_token = os.getenv('HUGGINGFACE_TOKEN')
 if not huggingface_token:
     raise ValueError("HUGGINGFACE_TOKEN environment variable is not set")
 
-model_id = "meta-llama/Llama-Guard-3-8B-INT8"
+model_id = "meta-llama/Llama-Guard-3-1B"
 device = "cuda" if torch.cuda.is_available() else "cpu"
 dtype = torch.bfloat16
 
-quantization_config = BitsAndBytesConfig(load_in_8bit=True)
-
 def parse_llama_guard_output(result):
     # "<END CONVERSATION>" 以降の部分を抽出
     safety_assessment = result.split("<END CONVERSATION>")[-1].strip()
@@ -43,7 +41,6 @@ def moderate(user_input, assistant_response):
         model_id, 
         torch_dtype=dtype, 
         device_map="auto",
-        quantization_config=quantization_config,
         token=huggingface_token,
         low_cpu_mem_usage=True
     )