Update app.py

app.py

@@ -26,11 +26,18 @@ def generate_text(
             return ""
         
         # トークナイザーとモデルのロード（GPUが使える場合はGPUへ移動）
-        tokenizer = AutoTokenizer.from_pretrained('Local-Novel-LLM-project/Vecteus-v1-abliterated', torch_dtype=torch.float16, attn_implementation="flash_attention_2", device_map="cuda", use_auth_token=HF_TOKEN)
+        tokenizer = AutoTokenizer.from_pretrained('Local-Novel-LLM-project/Vecteus-v1-abliterated', attn_implementation="flash_attention_2", use_auth_token=HF_TOKEN)
         model = AutoModelForCausalLM.from_pretrained('Local-Novel-LLM-project/Vecteus-v1-abliterated', use_auth_token=HF_TOKEN)
 
+        # GPUが利用可能ならGPUへ移動。bf16がサポートされている場合はbf16を使用
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+        if device == "cuda" and hasattr(torch.cuda, "is_bf16_supported") and torch.cuda.is_bf16_supported():
+            model.to(device, dtype=torch.bfloat16)
+        else:
+            model.to(device)
+
         # 入力テキストをトークン化してトークン数を取得
-        input_ids = tokenizer.encode(input_text, return_tensors="pt").to("cuda")
+        input_ids = tokenizer.encode(input_text, return_tensors="pt").to(device)
         input_token_count = input_ids.shape[1]
         
         # 総トークン数の上限を入力トークン数 + max_length（max_lengthはトークン数として扱う）