Update app.py

app.py

@@ -1,30 +1,47 @@
 import gradio as gr
-from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline, BitsAndBytesConfig
 import os
-import subprocess
+# import subprocess # 현재 코드에서 사용되지 않으므로 제거 가능
 import torch
 from huggingface_hub import login
 
-token = os.environ.get("HUGGING_FACE_HUB_TOKEN")
-login(token)
+# 환경 변수에서 토큰 가져오기
+token = os.environ.get("HF_TOKEN") # 일반적으로 "HF_TOKEN"으로 설정됩니다.
+if token:
+    login(token)
+else:
+    print("HF_TOKEN 환경 변수가 설정되지 않았습니다. 모델 다운로드에 문제가 있을 수 있습니다.")
 
 
-# ---------- STEP 1: Fine-tuned 모델 Git에서 clone ----------
-repo_url = "DMID23/MachineToolAgent"
+# ---------- STEP 1: Fine-tuned 모델 정보 ----------
+repo_id = "DMID23/MachineToolAgent" # 모델 저장소 ID
 
 
-# ---------- STEP 2: Tokenizer와 모델 로드 ----------
-#base_model = "mistralai/Mistral-7B-v0.1"
-#tokenizer = AutoTokenizer.from_pretrained(base_model)
+# ---------- STEP 2: 양자화 설정 및 모델 로드 ----------
+
+# 8bit 양자화 설정 (CPU 환경에서도 사용 가능)
+# load_in_8bit=True 옵션만으로도 BitsAndBytesConfig 객체를 자동으로 생성하여 적용합니다.
+# CPU에서는 float32 -> int8 양자화가 주로 일어납니다.
+quantization_config = BitsAndBytesConfig(load_in_8bit=True)
+
 model = AutoModelForCausalLM.from_pretrained(
-    repo_url,
-    torch_dtype=torch.float32,
+    repo_id,
+    quantization_config=quantization_config, # 양자화 설정 적용
+    torch_dtype=torch.float32, # 8비트 로드 시에도 내부적으로 float32로 처리되거나 혼합 정밀도로 작동할 수 있습니다.
+                               # 하지만 실제 메모리는 8비트만큼만 사용됩니다.
+    device_map="auto" # 모델의 각 레이어를 자동으로 최적의 장치(CPU/GPU)에 분배
+                      # CPU만 있다면 CPU로 로드됩니다.
 )
-model = model.to("cpu")
+print("Model loaded successfully.")
+
+# 만약 DMID23/MachineToolAgent 저장소에 토크나이저가 있다면 repo_id로 바꾸세요.
+tokenizer = AutoTokenizer.from_pretrained(repo_id)
 
+# pipe 설정 시, device=-1 (CPU) 명시
 pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, device=-1)
 
 # ---------- STEP 3: Gradio 함수 정의 ----------
+# (이 부분은 변경 없음)
 def generate_response(prompt, max_length=256, temperature=0.7):
     # max_length를 제한하여 속도를 빠르게 함
     outputs = pipe(
@@ -39,6 +56,7 @@ def generate_response(prompt, max_length=256, temperature=0.7):
     return outputs[0]["generated_text"]
 
 # ---------- STEP 4: Gradio UI ----------
+# (이 부분은 변경 없음)
 with gr.Blocks() as demo:
     gr.Markdown("# 🚀 Fine-tuned Mistral-7B (CPU Optimized)")
 
@@ -59,4 +77,4 @@ with gr.Blocks() as demo:
     )
     
 # ---------- STEP 5: Launch ----------
-demo.launch()
+demo.launch()