Spaces:
Sleeping
Sleeping
| """ | |
| Receipt LoRA Trainer Space | |
| Wraps finetune_receipt.py in a ZeroGPU Gradio app. | |
| Secrets required in this Space's settings: | |
| HF_TOKEN — your HF write token (huggingface.co/settings/tokens) | |
| """ | |
| import os | |
| import spaces | |
| import gradio as gr | |
| HF_REPO_ID = "summerdevlin46/llama-3.2-3b-receipt-lora" | |
| TRAINING_DATA = [ | |
| {"input": "MAHALAKSHMI MARKETING\nNo. 2816 Date: 27/5/26\nM/s. Veerabala (Mulal)\nParle 1 X 2450 = 2450\nBingo(C) 4 X 870 = 3480\nSubtotal 5930\nDiscount 612\nTotal 6542", "output": "{\"supplier\": \"Mahalakshmi Marketing\", \"invoice_no\": \"2816\", \"date\": \"2026-05-27\", \"items\": [{\"product_raw\": \"Parle\", \"qty_cases\": 1, \"qty_units\": 1, \"unit_cost\": 2450.0, \"total\": 2450.0}, {\"product_raw\": \"Bingo(C)\", \"qty_cases\": 4, \"qty_units\": 4, \"unit_cost\": 870.0, \"total\": 3480.0}], \"subtotal\": 5930.0, \"discount\": 612.0, \"gst\": 0.0, \"net_total\": 6542.0}"}, | |
| {"input": "SRI VENKATESHWARA MARKETING\nGSTIN: 36AZLIPV6442K12M\nCUSTOMER: VEERA BHADRA WS\nBill Date: 28/05/2026\n1 PARLE-G 100G QTY: 5/0 MRP: 10 SALE RATE: 8.625\n2 HAPPY 2 (24P)*13 QTY: 10/0 MRP: 9 SALE RATE: 4.464\nGROSS SALES: 8569.032\nSCHEMES: 168.352\nCASH DISC: 420.034\nGST: 210.017 SGST: 210.017\nNET AMOUNT: 8821.00", "output": "{\"supplier\": \"Sri Venkateshwara Marketing\", \"invoice_no\": \"SVM/26-27/2598\", \"date\": \"2026-05-28\", \"items\": [{\"product_raw\": \"PARLE-G 100G\", \"qty_cases\": 5, \"qty_units\": 120, \"unit_cost\": 8.625, \"total\": 1035.0}, {\"product_raw\": \"HAPPY 2 (24P)\", \"qty_cases\": 10, \"qty_units\": 240, \"unit_cost\": 4.464, \"total\": 1071.36}], \"subtotal\": 8569.032, \"discount\": 588.386, \"gst\": 420.034, \"net_total\": 8821.0}"}, | |
| {"input": "Brundhna Boys - 28/05\nhne 30X28 = 840\noam 50X9.5 450\nBm 10X9.5 95\nBm 5X12 50\nTotal 1435", "output": "{\"supplier\": \"sales_note\", \"invoice_no\": null, \"date\": \"2026-05-28\", \"items\": [{\"product_raw\": \"hne\", \"qty_cases\": 0, \"qty_units\": 30, \"unit_cost\": 28.0, \"total\": 840.0}, {\"product_raw\": \"oam\", \"qty_cases\": 0, \"qty_units\": 50, \"unit_cost\": 9.5, \"total\": 450.0}, {\"product_raw\": \"Bm\", \"qty_cases\": 0, \"qty_units\": 10, \"unit_cost\": 9.5, \"total\": 95.0}, {\"product_raw\": \"Bm\", \"qty_cases\": 0, \"qty_units\": 5, \"unit_cost\": 12.0, \"total\": 50.0}], \"subtotal\": 1435.0, \"discount\": 0.0, \"gst\": 0.0, \"net_total\": 1435.0}"}, | |
| {"input": "Mahalakshmi Marketing\nDocument Type: handwritten supplier bill\nBill No: 2816\nDate: 27/05/2026\n\nParle bulk 1 X 2450 = 2450\nBingo(C) 4 X 870 = 3480\nUnknown handwritten item 1 X 612 = 612\n\nTotal: 6542", "output": "{\"supplier\": \"Mahalakshmi Marketing\", \"invoice_no\": \"2816\", \"date\": \"2026-05-27\", \"items\": [{\"product_raw\": \"Parle bulk\", \"qty_cases\": 1, \"qty_units\": 1, \"unit_cost\": 2450.0, \"total\": 2450.0}, {\"product_raw\": \"Bingo(C)\", \"qty_cases\": 4, \"qty_units\": 4, \"unit_cost\": 870.0, \"total\": 3480.0}, {\"product_raw\": \"Unknown handwritten item\", \"qty_cases\": 0, \"qty_units\": 1, \"unit_cost\": 612.0, \"total\": 612.0, \"needs_review\": true}], \"subtotal\": 6542.0, \"discount\": 0.0, \"gst\": 0.0, \"net_total\": 6542.0}"}, | |
| {"input": "Brundavan Buns\nDocument Type: handwritten tally note\nDate: 28/05\n\nItem one 30 X 28 = 840\nOBM 50 X 9.5 = 475\nBun 10 X 9.5 = 95\nBun 5 X 10 = 50\n\nTotal: 1435", "output": "{\"supplier\": \"Brundavan Buns\", \"invoice_no\": null, \"date\": \"2026-05-28\", \"items\": [{\"product_raw\": \"Item one\", \"qty_cases\": 0, \"qty_units\": 30, \"unit_cost\": 28.0, \"total\": 840.0}, {\"product_raw\": \"OBM\", \"qty_cases\": 0, \"qty_units\": 50, \"unit_cost\": 9.5, \"total\": 475.0}, {\"product_raw\": \"Bun\", \"qty_cases\": 0, \"qty_units\": 10, \"unit_cost\": 9.5, \"total\": 95.0}, {\"product_raw\": \"Bun\", \"qty_cases\": 0, \"qty_units\": 5, \"unit_cost\": 10.0, \"total\": 50.0}], \"subtotal\": 1460.0, \"discount\": 25.0, \"gst\": 0.0, \"net_total\": 1435.0}"}, | |
| {"input": "Sri Venkateshwara Marketing\nDocument Type: printed tax invoice\nInvoice No: 6\nDate: 28/05/2026\n\nPARLE-G 60GM RS.72P | 5/0 | MRP 10.00 | RATE 8.625 | GST 5% | NET 3105.000\nHAPPY HAPPY 27.5G(24P)*13 | 10/0 | MRP 5.00 | RATE 4.464 | GST 5% | NET 5715.710\n\nGross Sales: 8759.032\nCGST: 210.017\nSGST: 210.017\nNet Amount: 8821.00", "output": "{\"supplier\": \"Sri Venkateshwara Marketing\", \"invoice_no\": \"6\", \"date\": \"2026-05-28\", \"items\": [{\"product_raw\": \"PARLE-G 60GM RS.72P\", \"qty_cases\": 5, \"qty_units\": 0, \"unit_cost\": 8.625, \"total\": 3105.0}, {\"product_raw\": \"HAPPY HAPPY 27.5G(24P)*13\", \"qty_cases\": 10, \"qty_units\": 0, \"unit_cost\": 4.464, \"total\": 5715.71}], \"subtotal\": 8759.032, \"discount\": 0.0, \"gst\": 420.034, \"net_total\": 8821.0}"}, | |
| ] | |
| SYSTEM_PROMPT = ( | |
| "You are a receipt parser for an Indian convenience store. " | |
| "Extract all line items from the receipt text. " | |
| "Return ONLY valid JSON, no markdown, no explanation." | |
| ) | |
| INSTRUCTION_TEMPLATE = """### Instruction: | |
| {system} | |
| ### Input: | |
| {input} | |
| ### Response: | |
| {output}""" | |
| def run_training() -> str: | |
| import traceback | |
| try: | |
| return _run_training_inner() | |
| except Exception: | |
| return traceback.format_exc() | |
| def _run_training_inner() -> str: | |
| import torch | |
| from unsloth import FastLanguageModel | |
| from trl import SFTTrainer | |
| from transformers import TrainingArguments | |
| from datasets import Dataset | |
| from huggingface_hub import HfApi, create_repo | |
| token = os.getenv("HF_TOKEN") | |
| if not token: | |
| return "ERROR: HF_TOKEN secret is not set. Add it in Space Settings → Repository secrets." | |
| log_lines = [] | |
| def log(msg: str) -> None: | |
| print(msg) | |
| log_lines.append(msg) | |
| log(f"GPU: {torch.cuda.get_device_name(0)}") | |
| log(f"Loading base model: unsloth/Llama-3.2-3B-Instruct-bnb-4bit") | |
| model, tokenizer = FastLanguageModel.from_pretrained( | |
| model_name="unsloth/Llama-3.2-3B-Instruct-bnb-4bit", | |
| max_seq_length=2048, | |
| dtype=None, | |
| load_in_4bit=True, | |
| ) | |
| model = FastLanguageModel.get_peft_model( | |
| model, | |
| r=16, | |
| target_modules=["q_proj", "k_proj", "v_proj", "o_proj", | |
| "gate_proj", "up_proj", "down_proj"], | |
| lora_alpha=16, | |
| lora_dropout=0.05, | |
| bias="none", | |
| use_gradient_checkpointing=True, | |
| ) | |
| records = [ | |
| {"text": INSTRUCTION_TEMPLATE.format( | |
| system=SYSTEM_PROMPT, | |
| input=ex["input"], | |
| output=ex["output"], | |
| )} | |
| for ex in TRAINING_DATA | |
| ] | |
| dataset = Dataset.from_list(records) | |
| log(f"Training on {len(dataset)} examples") | |
| trainer = SFTTrainer( | |
| model=model, | |
| tokenizer=tokenizer, | |
| train_dataset=dataset, | |
| dataset_text_field="text", | |
| max_seq_length=2048, | |
| args=TrainingArguments( | |
| per_device_train_batch_size=2, | |
| gradient_accumulation_steps=4, | |
| num_train_epochs=10, | |
| learning_rate=2e-4, | |
| fp16=True, | |
| logging_steps=1, | |
| output_dir="./llama-3.2-3b-receipt-lora", | |
| save_strategy="no", | |
| warmup_steps=5, | |
| optim="adamw_8bit", | |
| ), | |
| ) | |
| log("Training...") | |
| trainer.train() | |
| log("Training done.") | |
| log(f"Pushing adapter to Hub: {HF_REPO_ID}") | |
| create_repo(HF_REPO_ID, repo_type="model", exist_ok=True, token=token) | |
| model.push_to_hub(HF_REPO_ID, token=token) | |
| tokenizer.push_to_hub(HF_REPO_ID, token=token) | |
| log("Exporting GGUF (Q4_K_M)...") | |
| model.save_pretrained_gguf( | |
| "llama-3.2-3b-receipt", | |
| tokenizer, | |
| quantization_method="q4_k_m", | |
| ) | |
| gguf_filename = "llama-3.2-3b-receipt-unsloth.Q4_K_M.gguf" | |
| log(f"Uploading GGUF to Hub: {HF_REPO_ID}/{gguf_filename}") | |
| api = HfApi(token=token) | |
| api.upload_file( | |
| path_or_fileobj=gguf_filename, | |
| path_in_repo=gguf_filename, | |
| repo_id=HF_REPO_ID, | |
| repo_type="model", | |
| ) | |
| log(f"\nDone! Model published at: https://huggingface.co/{HF_REPO_ID}") | |
| log(f"Set in Dukaan Saathi .env: HF_RECEIPT_MODEL_REPO={HF_REPO_ID}") | |
| return "\n".join(log_lines) | |
| with gr.Blocks(title="Receipt LoRA Trainer") as demo: | |
| gr.Markdown( | |
| f""" | |
| # Receipt LoRA Trainer | |
| Fine-tunes **Llama-3.2-3B-Instruct** on 6 receipt text extraction examples using ZeroGPU (A100). | |
| Pushes the adapter + GGUF to **[{HF_REPO_ID}](https://huggingface.co/{HF_REPO_ID})**. | |
| **Before running:** add `HF_TOKEN` in Space Settings → Repository secrets. | |
| """ | |
| ) | |
| run_btn = gr.Button("Start fine-tuning", variant="primary") | |
| output = gr.Textbox(label="Training log", lines=30, max_lines=50) | |
| run_btn.click(fn=run_training, outputs=output) | |
| demo.launch() | |