Update app.py

app.py

@@ -6,7 +6,7 @@ import io
 from PIL import Image
 import torch
 from transformers import AutoProcessor, AutoModelForCausalLM
-from ctransformers import AutoModelForCausalLM as CAutoModelForCausalLM # Thư viện mới
+from ctransformers import AutoModelForCausalLM as CAutoModelForCausalLM
 
 # --- CẤU HÌNH VÀ TẢI MÔ HÌNH (PHIÊN BẢN ỔN ĐỊNH NHẤT CHO CPU) ---
 device = "cpu"
@@ -25,17 +25,15 @@ ocr_model = AutoModelForCausalLM.from_pretrained(
 )
 print("Tải xong mô hình OCR.")
 
-# 2. THAY ĐỔI LỚN: DÙNG CTRANSFORMERS ĐỂ TẢI GGUF
+# 2. DÙNG CTRANSFORMERS ĐỂ TẢI GGUF
 print("Đang tải mô hình LLM (Llama-3-8B GGUF for CPU)...")
-# Sử dụng phiên bản GGUF được tối ưu cho CPU
 llm_model_id = "QuantFactory/Meta-Llama-3-8B-Instruct-GGUF"
-llm_model_file = "Meta-Llama-3-8B-Instruct.Q4_K_M.gguf" # File quantization 4-bit, cân bằng tốt
+llm_model_file = "Meta-Llama-3-8B-Instruct.Q4_K_M.gguf"
 llm = CAutoModelForCausalLM.from_pretrained(
     llm_model_id,
     model_file=llm_model_file,
     model_type="llama",
-    # Cấu hình để chạy trên CPU
-    gpu_layers=0, # Không dùng GPU
+    gpu_layers=0,
     context_length=4096 
 )
 print("Tải xong mô hình LLM.")
@@ -47,6 +45,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)
+    # SỬA LỖI ĐÁNH MÁY Ở ĐÂY
     generated_ids = ocr_model.generate(
         input_ids=inputs["input_ids"],
         pixel_values=inputs["pixel_values"],
@@ -58,19 +57,14 @@ def run_ocr(image: Image.Image) -> str:
     return parsed_text['<OCR>']
 
 def extract_order_from_text(text: str) -> dict:
-    # Cập nhật prompt cho Llama 3 GGUF
     prompt = f"""<|begin_of_text|><|start_header_id|>system<|end_header_id|>
     You are an expert assistant that only outputs valid JSON. Extract order information from the text. The JSON object must contain "ten_khach_hang" (string, null if not found) and "danh_sach_hang" (an array of items). Each item must have "ten_hang" (string), "so_luong" (number), "don_vi" (string), "ma_hang" (string, null if not found), and "ghi_chu" (string, null if not found).<|eot_id|><|start_header_id|>user<|end_header_id|>
     Text Content:
     {text}<|eot_id|><|start_header_id|>assistant<|end_header_id|>
     """
-    
     response_text = llm(prompt, max_new_tokens=1024, temperature=0.1, stop=["<|eot_id|>"])
-    
     try:
-        # Làm sạch output thô từ GGUF model
         json_str = response_text.strip()
-        # Tìm vị trí bắt đầu và kết thúc của JSON
         start = json_str.find('{')
         end = json_str.rfind('}') + 1
         if start != -1 and end != 0:
@@ -81,17 +75,13 @@ def extract_order_from_text(text: str) -> dict:
 
 def create_excel_file(order_data: dict):
     if not order_data or "danh_sach_hang" not in order_data or not order_data["danh_sach_hang"]: return None
-    # Đổi tên key cho phù hợp với DataFrame
     flat_data = []
     customer = order_data.get('ten_khach_hang', 'N/A')
     for item in order_data['danh_sach_hang']:
         flat_data.append({
-            'Khách hàng': customer,
-            'Mã hàng': item.get('ma_hang'),
-            'Tên hàng': item.get('ten_hang'),
-            'Số lượng': item.get('so_luong'),
-            'Đơn vị': item.get('don_vi'),
-            'Ghi chú': item.get('ghi_chu')
+            'Khách hàng': customer, 'Mã hàng': item.get('ma_hang'),
+            'Tên hàng': item.get('ten_hang'), 'Số lượng': item.get('so_luong'),
+            'Đơn vị': item.get('don_vi'), 'Ghi chú': item.get('ghi_chu')
         })
     df = pd.DataFrame(flat_data)
     output = io.BytesIO()
@@ -121,7 +111,7 @@ def process_image_and_extract(image):
         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) ---
+# --- XÂY DỰNG GIAO DIỆN GRADIO ---
 with gr.Blocks(theme=gr.themes.Soft()) as app:
     gr.Markdown("# Ứng dụng Trích xuất Đơn hàng từ Ảnh chụp màn hình")
     gr.Markdown("Chụp màn hình email/tin nhắn đặt hàng, sau đó dán (Ctrl+V) vào ô bên dưới và nhấn 'Xử lý'.")