Update README.md

README.md

@@ -1,58 +1,182 @@
 ---
 base_model: HuggingFaceTB/SmolVLM2-500M-Video-Instruct
 library_name: transformers
-model_name: smolvlm2-500m-FoodExtract-Vision
+model_name: FoodExtract-Vision-SmolVLM2-500M-fine-tune
 tags:
-- generated_from_trainer
-- trl
+- vision-language-model
+- vlm
+- image-to-text
+- structured-output
+- food-extraction
+- fine-tuned
 - sft
-licence: license
+- trl
+license: apache-2.0
+datasets:
+- mrdbourke/FoodExtract-1k-Vision
+language:
+- en
+pipeline_tag: image-text-to-text
 ---
 
-# Model Card for smolvlm2-500m-FoodExtract-Vision
+# FoodExtract-Vision-SmolVLM2-500M
+
+A fine-tuned Vision-Language Model for **structured food and drink extraction** from images. Given an input image, the model outputs a structured JSON containing food classification, image title, and extracted food/drink items.
+
+## Model Description
+
+| Attribute | Value |
+|-----------|-------|
+| **Base Model** | [SmolVLM2-500M-Video-Instruct](https://huggingface.co/HuggingFaceTB/SmolVLM2-500M-Video-Instruct) |
+| **Training Method** | Supervised Fine-Tuning (SFT) |
+| **Training Strategy** | Vision Encoder Frozen, LLM & Cross-Modal Connector Trainable |
+| **Total Parameters** | 507M |
+| **Trainable Parameters** | 421M (83%) |
+| **Frozen Parameters** | 86M (17%) |
+| **Precision** | bfloat16 |
+
+## Intended Use
+
+This model is designed for:
+- 🍕 **Food/Drink Classification**: Determine if an image contains food or drinks
+- 📝 **Structured Data Extraction**: Extract food and drink items into JSON format
+- 🏷️ **Image Captioning**: Generate food-related titles for images
+
+### Output Format
 
-This model is a fine-tuned version of [HuggingFaceTB/SmolVLM2-500M-Video-Instruct](https://huggingface.co/HuggingFaceTB/SmolVLM2-500M-Video-Instruct).
-It has been trained using [TRL](https://github.com/huggingface/trl).
+```json
+{
+  "is_food": 1,
+  "image_title": "macaron assortment",
+  "food_items": ["yellow macaron", "white macaron", "green macaron"],
+  "drink_items": []
+}
+```
 
 ## Quick start
 
 ```python
 from transformers import pipeline
+import torch
+
+# Load the fine-tuned model
+pipe = pipeline(
+    "image-text-to-text",
+    model="CreatorJarvis/FoodExtract-Vision-SmolVLM2-500M-fine-tune",
+    dtype=torch.bfloat16,
+    device_map="auto"
+)
+
+# Prepare input message
+message = [{
+    "role": "user",
+    "content": [
+        {"type": "image", "image": your_image},  # PIL.Image object
+        {"type": "text", "text": """Classify the given input image into food or not and if edible food or drink items are present, extract those to a list. If no food/drink items are visible, return empty lists.
 
-question = "If you had a time machine, but could only go to the past or the future once and never return, which would you choose and why?"
-generator = pipeline("text-generation", model="None", device="cuda")
-output = generator([{"role": "user", "content": question}], max_new_tokens=128, return_full_text=False)[0]
-print(output["generated_text"])
+Only return valid JSON in the following form:
+
+```json
+{
+  'is_food': 0,
+  'image_title': '',
+  'food_items': [],
+  'drink_items': []
+}
+```} ] }]
 ```
 
-## Training procedure
 
-[<img src="https://raw.githubusercontent.com/wandb/assets/main/wandb-github-badge-28.svg" alt="Visualize in Weights & Biases" width="150" height="24"/>](https://wandb.ai/jarviszhang-new-york-university/vlm-food-extraction/runs/we6trzon) 
+## Training Details
+
+### Dataset
+
+| Split | Samples | Description |
+|-------|---------|-------------|
+| Train | 1,208 | 80% of total dataset |
+| Validation | 302 | 20% of total dataset |
+| **Total** | **1,510** | Food images (1k) + Non-food images (500) |
 
+**Dataset Source**: [mrdbourke/FoodExtract-1k-Vision](https://huggingface.co/datasets/mrdbourke/FoodExtract-1k-Vision)
 
-This model was trained with SFT.
+### Training Configuration
 
-### Framework versions
+| Hyperparameter | Value |
+|----------------|-------|
+| Epochs | 4 |
+| Batch Size (per device) | 4 |
+| Gradient Accumulation Steps | 4 |
+| Effective Batch Size | 16 |
+| Learning Rate | 2e-4 |
+| LR Scheduler | Constant |
+| Warmup Ratio | 0.03 |
+| Optimizer | AdamW (fused) |
+| Max Grad Norm | 1.0 |
+| Precision | bf16 |
+| Gradient Checkpointing | ✓ |
 
-- TRL: 0.27.1
-- Transformers: 4.57.6
-- Pytorch: 2.9.0+cu126
-- Datasets: 4.0.0
-- Tokenizers: 0.22.2
+### Training Strategy
 
-## Citations
+The **Vision Encoder was frozen** during training to:
+- Preserve pre-trained visual representations
+- Reduce trainable parameters and memory usage
+- Improve training stability on small datasets
+- Mitigate overfitting
 
+This approach is inspired by the [SmolDocling paper](https://arxiv.org/abs/2503.11576).
 
+### Training Results
+
+| Epoch | Training Loss | Validation Loss |
+|-------|---------------|-----------------|
+| 1 | 0.0842 | 0.0759 |
+| 2 | 0.0816 | 0.0757 |
+| 3 | 0.0237 | 0.0751 |
+| 4 | 0.0172 | 0.0807 |
+
+**Final Training Loss**: 0.0518
+
+### Experiment Tracking
+
+[<img src="https://raw.githubusercontent.com/wandb/assets/main/wandb-github-badge-28.svg" alt="Visualize in Weights & Biases" width="200" height="30"/>](https://wandb.ai/jarviszhang-new-york-university/vlm-food-extraction/runs/we6trzon)
+
+## Demo
+
+Try the model on Hugging Face Spaces:
+
+🚀 **[FoodExtract-Vision Demo](https://huggingface.co/spaces/CreatorJarvis/FoodExtract-Vision)**
+
+The demo compares outputs from the base model vs. the fine-tuned model side-by-side.
+
+## Limitations
+
+- Trained on a relatively small dataset (1.5k images)
+- May struggle with complex multi-item food scenes
+- Occasional repetitive generation patterns
+- Best performance on single-dish food images
+
+## Framework Versions
+
+| Library | Version |
+|---------|---------|
+| TRL | 0.27.1 |
+| Transformers | 4.57.6 |
+| PyTorch | 2.9.0+cu126 |
+| Datasets | 4.0.0 |
+| Tokenizers | 0.22.2 |
+
+## Citation
+
+If you use this model, please cite:
 
-Cite TRL as:
-    
 ```bibtex
-@misc{vonwerra2022trl,
-	title        = {{TRL: Transformer Reinforcement Learning}},
-	author       = {Leandro von Werra and Younes Belkada and Lewis Tunstall and Edward Beeching and Tristan Thrush and Nathan Lambert and Shengyi Huang and Kashif Rasul and Quentin Gallou{\'e}dec},
-	year         = 2020,
-	journal      = {GitHub repository},
-	publisher    = {GitHub},
-	howpublished = {\url{https://github.com/huggingface/trl}}
+@misc{foodextract-vision-2025,
+  title        = {FoodExtract-Vision: Fine-tuned SmolVLM2 for Structured Food Extraction},
+  author       = {Jarvis Zhang},
+  year         = 2025,
+  publisher    = {Hugging Face},
+  howpublished = {\url{[https://huggingface.co/CreatorJarvis/FoodExtract-Vision-SmolVLM2-500M-fine-tune](https://huggingface.co/CreatorJarvis/FoodExtract-Vision-SmolVLM2-500M-fine-tune)}}
 }
+
+
 ```