Update app.py

app.py

@@ -5,35 +5,33 @@ import json
 import io
 from PIL import Image
 import torch
-from transformers import AutoProcessor, AutoModelForCausalLM, BitsAndBytesConfig, pipeline
+from transformers import AutoProcessor, AutoModelForCausalLM, pipeline
 
-# --- CẤU HÌNH VÀ TẢI MÔ HÌNH (PHIÊN BẢN NHẸ HƠN) ---
-device = "cuda" if torch.cuda.is_available() else "cpu"
+# --- CẤU HÌNH VÀ TẢI MÔ HÌNH (PHIÊN BẢN ỔN ĐỊNH CHO CPU) ---
+device = "cpu" # Ép chạy trên CPU để đảm bảo
 print(f"Đang sử dụng thiết bị: {device}")
 
-# 1. Tải mô hình OCR (Sử dụng phiên bản 'base' thay vì 'large')
+# 1. Tải mô hình OCR (Không nén)
 print("Đang tải mô hình OCR (Florence-2-base)...")
 ocr_model_id = "microsoft/Florence-2-base"
 ocr_processor = AutoProcessor.from_pretrained(ocr_model_id, trust_remote_code=True)
-quantization_config = BitsAndBytesConfig(load_in_4bit=True)
 ocr_model = AutoModelForCausalLM.from_pretrained(
     ocr_model_id, 
-    device_map="auto", 
-    torch_dtype=torch.bfloat16,
-    quantization_config=quantization_config,
+    device_map=device, # Chạy trên CPU
+    torch_dtype=torch.float32, # Dùng float32 cho CPU
     trust_remote_code=True,
-    attn_implementation="eager" # <-- THÊM DÒNG NÀY ĐỂ SỬA LỖI
+    attn_implementation="eager" 
 )
 print("Tải xong mô hình OCR.")
 
-# 2. Tải mô hình LLM (Sử dụng phiên bản đã được nén sẵn)
-print("Đang tải mô hình LLM (Unsloth Llama 3 8B 4-bit)...")
-llm_model_id = "unsloth/llama-3-8b-Instruct-bnb-4bit"
+# 2. Tải mô hình LLM (Không nén)
+print("Đang tải mô hình LLM (Meta Llama 3 8B)...")
+llm_model_id = "meta-llama/Meta-Llama-3-8B-Instruct" # Dùng phiên bản gốc
 llm_pipeline = pipeline(
     "text-generation",
     model=llm_model_id,
-    model_kwargs={"torch_dtype": torch.bfloat16},
-    device_map="auto",
+    model_kwargs={"torch_dtype": torch.float32}, # Dùng float32 cho CPU
+    device=0 if device == "cuda" else -1, # -1 để pipeline dùng CPU
 )
 print("Tải xong mô hình LLM.")
 
@@ -44,7 +42,7 @@ def run_ocr(image: Image.Image) -> str:
     if image is None: return "Lỗi: Vui lòng cung cấp hình ảnh."
     prompt = "<OCR>"
     inputs = ocr_processor(text=prompt, images=image, return_tensors="pt").to(device)
-    generated_ids = ocr_model.generate(
+    generated_ids = oocr_model.generate(
         input_ids=inputs["input_ids"],
         pixel_values=inputs["pixel_values"],
         max_new_tokens=2048,
@@ -74,18 +72,35 @@ def create_excel_file(order_data: dict):
     return (filename, output.getvalue())
 
 def process_image_and_extract(image):
-    if image is None: return "Vui lòng dán ảnh vào.", None, None
-    extracted_text = run_ocr(image)
-    if not extracted_text.strip(): return "Không đọc được chữ từ hình ảnh.", None, None
-    order_data = extract_order_from_text(extracted_text)
-    if "error" in order_data: return extracted_text, f"Lỗi từ AI: {order_data['error']}\nPhản hồi gốc: {order_data['raw_response']}", None
-    excel_info = create_excel_file(order_data)
-    df_display = pd.DataFrame(order_data.get('danh_sach_hang', []))
-    if excel_info:
-        filename, filebytes = excel_info
-        with open(filename, "wb") as f: f.write(filebytes)
-        return extracted_text, df_display, filename
-    else: return extracted_text, df_display, None
+    try:
+        if image is None: return "Vui lòng dán ảnh vào.", None, None
+        
+        print("Bắt đầu OCR...")
+        extracted_text = run_ocr(image)
+        if not extracted_text.strip(): return "Không đọc được chữ từ hình ảnh.", None, None
+        print(f"Văn bản OCR: {extracted_text}")
+
+        print("Bắt đầu trích xuất LLM...")
+        order_data = extract_order_from_text(extracted_text)
+        if "error" in order_data: return extracted_text, f"Lỗi từ AI: {order_data['error']}\nPhản hồi gốc: {order_data['raw_response']}", None
+        print(f"Dữ liệu trích xuất: {order_data}")
+
+        excel_info = create_excel_file(order_data)
+        df_display = pd.DataFrame(order_data.get('danh_sach_hang', []))
+        
+        if excel_info:
+            filename, filebytes = excel_info
+            with open(filename, "wb") as f: f.write(filebytes)
+            return extracted_text, df_display, filename
+        else: return extracted_text, df_display, None
+    except Exception as e:
+        # Bắt lỗi và hiển thị trong giao diện để dễ gỡ lỗi
+        import traceback
+        error_str = str(e)
+        traceback_str = traceback.format_exc()
+        print(traceback_str)
+        return f"Lỗi nghiêm trọng: {error_str}", None, None
+
 
 # --- XÂY DỰNG GIAO DIỆN GRADIO (GIỮ NGUYÊN) ---