Update app.py

app.py

@@ -9,11 +9,10 @@ Pipeline:
     → deterministic complete()  (VAT · serial · allocation number)
     → fiscal document display
     → FAISS recommender  (3 similar past receipts)
-
 Repos used:
   Dataset : yonilev/Text2Receipt
-  Model   : yonilev/Text2Receipt-parser  (LoRA adapter + embeddings)
-  Space   : yonilev/Text2Receipt
+  Model   : yonilev/Text2Receipt-parser  (LoRA adapter)
+  Space   : yonilev/Text2Receipt         (embeddings stored here)
 """
 import os, json, re, random
 import numpy as np
@@ -39,54 +38,71 @@ def _lazy_init():
     from peft import PeftModel
     from sentence_transformers import SentenceTransformer
     import faiss
-    from huggingface_hub import hf_hub_download
     import pandas as pd
 
     device = "cuda" if torch.cuda.is_available() else "cpu"
     _STATE["device"] = device
     _STATE["rng"]    = random.Random(42)
 
-    # ── parser model ──────────────────────────────────────────────────────────
-    bnb = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_quant_type="nf4",
-                             bnb_4bit_compute_dtype=torch.float16,
-                             bnb_4bit_use_double_quant=True)
+    # ── parser model — CPU-safe loading ──────────────────────────────────────
     tok = AutoTokenizer.from_pretrained(BASE_MODEL)
     if tok.pad_token is None:
         tok.pad_token = tok.eos_token
     tok.padding_side = "left"
 
-    try:
-        base = AutoModelForCausalLM.from_pretrained(
-            BASE_MODEL, quantization_config=bnb, device_map="auto",
-            torch_dtype=torch.float16, attn_implementation="eager")
-        model = PeftModel.from_pretrained(base, MODEL_REPO)
-        print("✅ loaded fine-tuned adapter from", MODEL_REPO)
-    except Exception as e:
-        print(f"⚠ adapter load failed ({e}); falling back to base model")
-        model = AutoModelForCausalLM.from_pretrained(
-            BASE_MODEL, quantization_config=bnb, device_map="auto",
-            torch_dtype=torch.float16, attn_implementation="eager")
+    if device == "cuda":
+        # GPU path: 4-bit quantization
+        bnb = BitsAndBytesConfig(
+            load_in_4bit=True,
+            bnb_4bit_quant_type="nf4",
+            bnb_4bit_compute_dtype=torch.float16,
+            bnb_4bit_use_double_quant=True,
+        )
+        try:
+            base = AutoModelForCausalLM.from_pretrained(
+                BASE_MODEL, quantization_config=bnb, device_map="auto",
+                torch_dtype=torch.float16, attn_implementation="eager")
+            model = PeftModel.from_pretrained(base, MODEL_REPO)
+            print("✅ GPU: loaded fine-tuned adapter from", MODEL_REPO)
+        except Exception as e:
+            print(f"⚠ adapter load failed ({e}); falling back to base model (GPU)")
+            model = AutoModelForCausalLM.from_pretrained(
+                BASE_MODEL, quantization_config=bnb, device_map="auto",
+                torch_dtype=torch.float16, attn_implementation="eager")
+    else:
+        # CPU path: fp32, no quantization (bitsandbytes not supported on CPU)
+        try:
+            base = AutoModelForCausalLM.from_pretrained(
+                BASE_MODEL, device_map="cpu",
+                torch_dtype=torch.float32, attn_implementation="eager")
+            model = PeftModel.from_pretrained(base, MODEL_REPO)
+            print("✅ CPU: loaded fine-tuned adapter from", MODEL_REPO)
+        except Exception as e:
+            print(f"⚠ adapter load failed ({e}); falling back to base model (CPU)")
+            model = AutoModelForCausalLM.from_pretrained(
+                BASE_MODEL, device_map="cpu",
+                torch_dtype=torch.float32, attn_implementation="eager")
 
     model.eval()
     _STATE["tok"]   = tok
     _STATE["model"] = model
 
-    # ── embeddings + FAISS ────────────────────────────────────────────────────
-    emb_path   = hf_hub_download(MODEL_REPO, "receipts_embeddings.npy", repo_type="model")
-    store_path = hf_hub_download(MODEL_REPO, "receipts_store.parquet",  repo_type="model")
-    manifest_path = hf_hub_download(MODEL_REPO, "embeddings_manifest.json", repo_type="model")
+    # ── embeddings + FAISS — load from Space files (co-located) ─────────────
+    # Files are uploaded directly to the Space repo, so they live at ./
+    SPACE_EMB_PATH      = "receipts_embeddings.npy"
+    SPACE_STORE_PATH    = "receipts_store.parquet"
+    SPACE_MANIFEST_PATH = "embeddings_manifest.json"
 
-    emb   = np.load(emb_path).astype("float32")
-    store = pd.read_parquet(store_path)
-    with open(manifest_path) as f:
+    emb   = np.load(SPACE_EMB_PATH).astype("float32")
+    store = pd.read_parquet(SPACE_STORE_PATH)
+    with open(SPACE_MANIFEST_PATH) as f:
         manifest = json.load(f)
 
     faiss.normalize_L2(emb)
     index = faiss.IndexFlatIP(emb.shape[1])
     index.add(emb)
 
-    enc = SentenceTransformer(manifest["embed_model"],
-                              device=device)
+    enc = SentenceTransformer(manifest["embed_model"], device=device)
     _STATE["enc"]       = enc
     _STATE["index"]     = index
     _STATE["store"]     = store
@@ -149,7 +165,6 @@ def _today() -> _dt.date:
     return _dt.date.today()
 
 def _missing_fields(parse: dict | None) -> list[str]:
-    """Return list of missing mandatory fields that affect the document legally."""
     missing = []
     if parse is None:
         return ["parse_failed"]
@@ -157,11 +172,9 @@ def _missing_fields(parse: dict | None) -> list[str]:
         missing.append("client_name")
     if not parse.get("items"):
         missing.append("items")
-    # doc_type not extractable from notes — always needs clarification for allocation
     return missing
 
 def _needs_allocation_clarification(parse: dict) -> bool:
-    """Return True when the doc_type choice changes whether allocation is required."""
     if parse is None:
         return False
     subtotal = sum(
@@ -170,15 +183,9 @@ def _needs_allocation_clarification(parse: dict) -> bool:
     )
     threshold = core.allocation_threshold_for_date(_today())
     is_biz    = parse.get("client_is_business", False)
-    # allocation only matters for authorized dealer + biz client + subtotal ≥ threshold
-    # but we don't know issuer status here; ask when subtotal is near or above threshold
-    return bool(is_biz) and subtotal >= threshold * 0.8  # 80% buffer
+    return bool(is_biz) and subtotal >= threshold * 0.8
 
 def clarification_questions(parse: dict | None, answers: dict) -> list[str]:
-    """
-    Agentic loop: generate next question(s) given current parse + prior answers.
-    Returns [] when all necessary info is available.
-    """
     questions = []
     missing = _missing_fields(parse)
 
@@ -191,7 +198,6 @@ def clarification_questions(parse: dict | None, answers: dict) -> list[str]:
     if "items" in missing and "items" not in answers:
         questions.append("מה השירות/המוצר שסופק ובאיזה מחיר?")
 
-    # ask doc_type only if not already answered and allocation is relevant
     if "doc_type" not in answers:
         if parse and _needs_allocation_clarification(parse):
             questions.append(
@@ -202,11 +208,10 @@ def clarification_questions(parse: dict | None, answers: dict) -> list[str]:
                 "איזה מסמך לייצר? קבלה / חשבונית מס / חשבונית מס וקבלה?"
             )
 
-    return questions[:3]  # max 3 questions per round
+    return questions[:3]
 
 
 def apply_answers(parse: dict | None, answers: dict) -> dict:
-    """Merge user answers into parse dict."""
     if parse is None:
         parse = {}
     p = dict(parse)
@@ -220,7 +225,6 @@ def apply_answers(parse: dict | None, answers: dict) -> dict:
             p["doc_type"] = "tax_invoice"
         else:
             p["doc_type"] = "receipt"
-    # set safe defaults for missing required fields
     p.setdefault("doc_type",        "receipt")
     p.setdefault("date",            _today().isoformat())
     p.setdefault("payment_method",  "bank_transfer")
@@ -236,14 +240,13 @@ def apply_answers(parse: dict | None, answers: dict) -> dict:
 # 4.  FAISS recommender
 # ═════════════════════════════════════════════════════════════════════════════
 def recommend(query_text: str, k: int = 3) -> list[dict]:
-    """Return k most similar receipts from the full corpus store."""
     enc   = _STATE["enc"]
     index = _STATE["index"]
     store = _STATE["store"]
     pref  = "query: " if _STATE["e5_family"] else ""
     q_emb = enc.encode([pref + query_text],
                        normalize_embeddings=True).astype("float32")
-    _, I = index.search(q_emb, k + 1)  # +1 in case query is in corpus
+    _, I = index.search(q_emb, k + 1)
     results = []
     for idx in I[0][:k]:
         row = store.iloc[int(idx)]
@@ -288,8 +291,6 @@ def render_document(completed: dict) -> str:
      background:#0f1117; color:#e8eaf6; border-radius:16px;
      padding:28px 32px; max-width:680px; margin:0 auto;
      border:1px solid #2a2d3e; box-shadow:0 4px 32px rgba(0,0,0,.4);">
-
-  <!-- header -->
   <div style="display:flex;justify-content:space-between;align-items:flex-start;
        border-bottom:2px solid #0d9488;padding-bottom:16px;margin-bottom:20px;">
     <div>
@@ -308,8 +309,6 @@ def render_document(completed: dict) -> str:
       </div>
     </div>
   </div>
-
-  <!-- client -->
   <div style="background:#1a1d2e;border-radius:10px;padding:14px 18px;
        margin-bottom:20px;font-size:14px;">
     <span style="color:#9e9e9e;">לקוח: </span>
@@ -317,10 +316,7 @@ def render_document(completed: dict) -> str:
     {'<span style="margin-right:10px;font-size:12px;color:#9e9e9e;">עסק</span>' if client.get('is_business') else ''}
     {f'<span style="font-size:12px;color:#9e9e9e;"> · ח.פ. {client.get("tax_id","")}</span>' if client.get('tax_id') else ''}
   </div>
-
   {alloc_html}
-
-  <!-- line items -->
   <table style="width:100%;border-collapse:collapse;font-size:14px;
          margin-bottom:18px;">
     <thead>
@@ -335,8 +331,6 @@ def render_document(completed: dict) -> str:
       {line_rows}
     </tbody>
   </table>
-
-  <!-- totals -->
   <div style="border-top:1px solid #2a2d3e;padding-top:14px;
        font-size:14px;text-align:left;">
     <div style="display:flex;justify-content:space-between;margin-bottom:4px;">
@@ -438,19 +432,13 @@ DEMO_ISSUER = {
 
 def _run_pipeline(raw_text: str, chat_history: list,
                   pending_parse: dict | None, answers: dict):
-    """
-    Main pipeline step.
-    Returns: (chat_history, doc_html, recs_html, pending_parse, answers)
-    """
     _lazy_init()
 
-    # Step 1 — parse
     if pending_parse is None:
         parse = model_parse(raw_text)
     else:
         parse = pending_parse
 
-    # Step 2 — check for questions
     questions = clarification_questions(parse, answers)
     if questions:
         q_text = "\n".join(f"• {q}" for q in questions)
@@ -460,7 +448,6 @@ def _run_pipeline(raw_text: str, chat_history: list,
         ]
         return chat_history, "", "", parse, answers
 
-    # Step 3 — apply answers + complete
     final_parse = apply_answers(parse, answers)
     try:
         completed = core.complete(DEMO_ISSUER, final_parse, _STATE["rng"])
@@ -469,7 +456,6 @@ def _run_pipeline(raw_text: str, chat_history: list,
                                  "content": f"שגיאה בעיבוד: {e}"}],
                 "", "", None, {})
 
-    # Step 4 — render
     doc_html  = render_document(completed)
     recs      = recommend(raw_text)
     recs_html = render_recommendations(recs)
@@ -484,11 +470,9 @@ def _run_pipeline(raw_text: str, chat_history: list,
 def _handle_answer(user_msg: str, chat_history: list,
                    pending_parse: dict | None, answers: dict,
                    raw_text: str):
-    """Parse a free-text answer and route back through pipeline."""
     a = dict(answers)
     low = user_msg.strip()
 
-    # simple intent detection for clarification answers
     if any(w in low for w in ["חשבונית מס", "מס וקבלה", "מס/קבלה"]):
         a["doc_type"] = low
     elif "קבלה" in low:
@@ -501,7 +485,7 @@ def _handle_answer(user_msg: str, chat_history: list,
 
 
 # ── build UI ─────────────────────────────────────────────────────────────────
-with gr.Blocks(title="Text2Receipt") as demo:
+with gr.Blocks(title="Text2Receipt", css=CSS) as demo:
 
     gr.HTML("""
     <div dir="rtl" style="text-align:center;padding:24px 0 8px;">
@@ -512,7 +496,6 @@ with gr.Blocks(title="Text2Receipt") as demo:
       </div>
     </div>""")
 
-    # state
     st_parse   = gr.State(None)
     st_answers = gr.State({})
     st_raw     = gr.State("")
@@ -522,7 +505,8 @@ with gr.Blocks(title="Text2Receipt") as demo:
             note_input = gr.Textbox(
                 label="הערת הכנסה (עברית חופשית)",
                 placeholder="לדוגמה: קיבלתי 1,200 ש\"ח ממשה כהן על ייעוץ עסקי",
-                lines=3,            )
+                lines=3,
+            )
             with gr.Row():
                 submit_btn = gr.Button("⚡ הפק מסמך", variant="primary")
                 clear_btn  = gr.Button("🗑 נקה")
@@ -536,17 +520,18 @@ with gr.Blocks(title="Text2Receipt") as demo:
 
         with gr.Column(scale=3):
             chatbot = gr.Chatbot(
-                label="סוכן הבהרה",                height=220,
+                label="סוכן הבהרה",
+                height=220,
             )
             answer_input = gr.Textbox(
                 label="תשובה לשאלת הסוכן",
                 placeholder="הקלד תשובה ולחץ Enter...",
-                visible=True,            )
+                visible=True,
+            )
 
     doc_output  = gr.HTML(label="מסמך פיסקלי")
     recs_output = gr.HTML(label="קבלות דומות")
 
-    # ── event wiring ──────────────────────────────────────────────────────────
     def on_submit(raw, history, pending, answers):
         history = history + [{"role": "user", "content": raw}]
         return _run_pipeline(raw, history, None, {})
@@ -572,4 +557,4 @@ with gr.Blocks(title="Text2Receipt") as demo:
     )
 
 if __name__ == "__main__":
-    demo.launch(css=CSS)
+    demo.launch(css=CSS)