Update app.py

app.py

@@ -5,31 +5,30 @@ import json
 import io
 from PIL import Image
 import torch
-from transformers import AutoProcessor, AutoModelForCausalLM, BitsAndBytesConfig
+from transformers import AutoProcessor, AutoModelForCausalLM, BitsAndBytesConfig, pipeline
 
-# --- CẤU HÌNH VÀ TẢI MÔ HÌNH ---
-# Sử dụng GPU nếu có, nếu không thì dùng CPU
+# --- 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"
 print(f"Đang sử dụng thiết bị: {device}")
 
-# 1. Tải mô hình OCR (Optical Character Recognition)
-print("Đang tải mô hình OCR (Florence-2)...")
-ocr_model_id = "microsoft/Florence-2-large"
+# 1. Tải mô hình OCR (Sử dụng phiên bản 'base' thay vì 'large')
+print("Đang tải mô hình OCR (Florence-2-base)...")
+ocr_model_id = "microsoft/Florence-2-base" # THAY ĐỔI 1: Dùng model base nhỏ hơn
 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=device, 
+    device_map="auto", # Để transformers tự quyết định device map
     torch_dtype=torch.bfloat16,
     quantization_config=quantization_config,
     trust_remote_code=True
 )
 print("Tải xong mô hình OCR.")
 
-# 2. Tải mô hình LLM (Language Model)
-print("Đang tải mô hình LLM (Meta Llama 3 8B)...")
-llm_model_id = "meta-llama/Meta-Llama-3-8B-Instruct"
-from transformers import pipeline
+# 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)...")
+# THAY ĐỔI 2: Dùng phiên bản Llama-3 đã được nén sẵn của Unsloth, rất hiệu quả về bộ nhớ
+llm_model_id = "unsloth/llama-3-8b-Instruct-bnb-4bit"
 llm_pipeline = pipeline(
     "text-generation",
     model=llm_model_id,
@@ -38,10 +37,10 @@ llm_pipeline = pipeline(
 )
 print("Tải xong mô hình LLM.")
 
-# --- CÁC HÀM XỬ LÝ ---
+
+# --- CÁC HÀM XỬ LÝ (GIỮ NGUYÊN) ---
 
 def run_ocr(image: Image.Image) -> str:
-    """Hàm chạy OCR để đọc chữ từ ảnh"""
     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)
@@ -56,7 +55,6 @@ def run_ocr(image: Image.Image) -> str:
     return parsed_text['<OCR>']
 
 def extract_order_from_text(text: str) -> dict:
-    """Hàm chạy LLM để trích xuất thông tin từ văn bản"""
     prompt = f"""You are an expert assistant for extracting order information from unstructured text. Based on the following text, extract the information and return it as a valid JSON object. The JSON object must contain: "ten_khach_hang": The customer's name (null if not found). "danh_sach_hang": An array of items. Each item must have: "ten_hang", "so_luong" (as a number), "don_vi", "ma_hang" (null if not found), and "ghi_chu" (null if not found). Output only the JSON object, with no additional text or explanation. --- Text Content --- {text} --- End Text Content ---"""
     messages = [{"role": "system", "content": "You are an assistant that only outputs valid JSON."}, {"role": "user", "content": prompt},]
     terminators = [llm_pipeline.tokenizer.eos_token_id, llm_pipeline.tokenizer.convert_tokens_to_ids("<|eot_id|>")]
@@ -66,7 +64,6 @@ def extract_order_from_text(text: str) -> dict:
     except json.JSONDecodeError: return {"error": "AI trả về định dạng không hợp lệ", "raw_response": response_text}
 
 def create_excel_file(order_data: dict):
-    """Hàm tạo file Excel từ dữ liệu đã trích xuất"""
     if not order_data or "danh_sach_hang" not in order_data or not order_data["danh_sach_hang"]: return None
     flat_data = [{'Khách hàng': order_data.get('ten_khach_hang', 'N/A'), **item} for item in order_data['danh_sach_hang']]
     df = pd.DataFrame(flat_data)
@@ -77,7 +74,6 @@ def create_excel_file(order_data: dict):
     return (filename, output.getvalue())
 
 def process_image_and_extract(image):
-    """Hàm tổng hợp, gọi OCR rồi đến LLM"""
     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
@@ -91,7 +87,7 @@ def process_image_and_extract(image):
         return extracted_text, df_display, filename
     else: return extracted_text, df_display, None
 
-# --- XÂY DỰNG GIAO DIỆN GRADIO ---
+# --- XÂY DỰNG GIAO DIỆN GRADIO (GIỮ NGUYÊN) ---
 
 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")