Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,4 +1,10 @@
 import streamlit as st
+st.set_page_config(
+    page_title="🍽️ Food Nutrition Estimator",
+    page_icon="🥗",
+    layout="centered"
+)
+
 import torch
 import os
 from PIL import Image
@@ -9,130 +15,124 @@ from transformers import (
     AutoModelForCausalLM
 )
 
-# 1. Streamlit UI setup
-st.set_page_config(
-    page_title="🍽️ Food Nutrition Estimator",
-    page_icon="🥗",
-    layout="centered"
-)
-st.title("🍽️ Food Nutrition Estimator")
-st.markdown("Upload a food image and get a nutritional overview generated by an instruction‐tuned LLM!")
-
-# 2. Environment & cache
-hf_token = os.getenv("HF_TOKEN", None)
-cache_dir = "/tmp/cache"
-os.makedirs(cache_dir, exist_ok=True)
-os.environ["HUGGINGFACE_HUB_CACHE"] = cache_dir
-
-# 3. Image transform for ViT
-manual_transform = transforms.Compose([
-    transforms.Resize(256),
-    transforms.CenterCrop(224),
-    transforms.Lambda(lambda img: img.convert("RGB")),
-    transforms.ToTensor(),
-    transforms.Normalize(mean=[0.485, 0.456, 0.406],
-                         std=[0.229, 0.224, 0.225]),
-    transforms.ConvertImageDtype(torch.float32)
-])
-
-# 4. Sidebar info
-st.sidebar.header("Models Used")
-st.sidebar.markdown("""
-- 🖼️ **Image Classifier**: `shingguy1/fine_tuned_vit`  
-- 💬 **Text Generator**: `tiiuae/falcon-7b-instruct`
-""")
-
-# 5. Load models (cached)
-@st.cache_resource
-def load_models():
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
-    # ViT classifier
-    model_vit = ViTForImageClassification.from_pretrained(
-        "shingguy1/fine_tuned_vit",
-        cache_dir=cache_dir,
-        use_auth_token=hf_token
-    ).to(device)
-
-    # Falcon‐7B Instruct LLM
-    tokenizer_llm = AutoTokenizer.from_pretrained(
-        "tiiuae/falcon-7b-instruct",
-        cache_dir=cache_dir,
-        use_auth_token=hf_token
-    )
-    model_llm = AutoModelForCausalLM.from_pretrained(
-        "tiiuae/falcon-7b-instruct",
-        cache_dir=cache_dir,
-        use_auth_token=hf_token,
-        torch_dtype=torch.float16,
-        device_map="auto"
-    )
-
-    return model_vit, tokenizer_llm, model_llm, device
-
-model_vit, tokenizer_llm, model_llm, device = load_models()
-
-# 6. Image uploader
-uploaded_file = st.file_uploader("Upload a food image...", type=["jpg", "jpeg", "png"])
-
-if uploaded_file is not None:
-    try:
-        # Display image
-        image = Image.open(uploaded_file)
-        st.image(image, caption="Uploaded Image", use_column_width=True)
-
-        # Classify with ViT
-        input_tensor = manual_transform(image).unsqueeze(0).to(device)
-        with torch.no_grad():
-            outputs = model_vit(pixel_values=input_tensor)
-        pred_idx = outputs.logits.argmax(-1).item()
-        pred_label = model_vit.config.id2label[pred_idx]
-        st.success(f"🍴 Predicted Food: **{pred_label}**")
-
-        # Build a single, unified instruction prompt
-        prompt = (
-            "### Instruction\n"
-            f"Provide a concise nutritional overview for a {pred_label}, including:\n"
-            "- Serving size (exact measurements & ingestion guidelines)\n"
-            "- Calories\n"
-            "- Protein, carbohydrates, and fat\n"
-            "- Main ingredients\n"
-            "- Cooking method\n"
-            "- One healthy substitution\n"
-            "### Response"
+def main():
+    # 2. Environment & cache
+    hf_token = os.getenv("HF_TOKEN", None)
+    cache_dir = "/tmp/cache"
+    os.makedirs(cache_dir, exist_ok=True)
+    os.environ["HUGGINGFACE_HUB_CACHE"] = cache_dir
+
+    # 3. Image transform for ViT
+    manual_transform = transforms.Compose([
+        transforms.Resize(256),
+        transforms.CenterCrop(224),
+        transforms.Lambda(lambda img: img.convert("RGB")),
+        transforms.ToTensor(),
+        transforms.Normalize(mean=[0.485, 0.456, 0.406],
+                             std=[0.229, 0.224, 0.225]),
+        transforms.ConvertImageDtype(torch.float32)
+    ])
+
+    # 4. Sidebar info
+    st.sidebar.header("Models Used")
+    st.sidebar.markdown("""
+    - 🖼️ **Image Classifier**: `shingguy1/fine_tuned_vit`  
+    - 💬 **Text Generator**: `tiiuae/falcon-7b-instruct`
+    """)
+
+    # 5. Load models
+    @st.cache_resource
+    def load_models():
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+        # ViT classifier
+        model_vit = ViTForImageClassification.from_pretrained(
+            "shingguy1/fine_tuned_vit",
+            cache_dir=cache_dir,
+            use_auth_token=hf_token
+        ).to(device)
+
+        # Falcon-7B Instruct LLM
+        tokenizer_llm = AutoTokenizer.from_pretrained(
+            "tiiuae/falcon-7b-instruct",
+            cache_dir=cache_dir,
+            use_auth_token=hf_token
         )
-        st.subheader("🧾 Nutrition Information")
-        st.write(f"🤖 Prompt sent to LLM:\n\n{prompt}")
-
-        # Tokenize & generate
-        inputs = tokenizer_llm(prompt, return_tensors="pt")
-        inputs = {k: v.to(model_llm.device) for k, v in inputs.items()}
-        input_len = inputs["input_ids"].shape[1]
-
-        outputs = model_llm.generate(
-            **inputs,
-            max_length=input_len + 150,
-            temperature=0.7,
-            top_p=0.9,
-            do_sample=True,
-            no_repeat_ngram_size=2,
-            early_stopping=True,
-            pad_token_id=tokenizer_llm.eos_token_id,
-            eos_token_id=tokenizer_llm.eos_token_id
+        model_llm = AutoModelForCausalLM.from_pretrained(
+            "tiiuae/falcon-7b-instruct",
+            cache_dir=cache_dir,
+            use_auth_token=hf_token,
+            torch_dtype=torch.float16,
+            device_map="auto"
         )
 
-        # Decode and strip prompt
-        full = tokenizer_llm.decode(outputs[0], skip_special_tokens=True).strip()
-        if full.startswith("### Response"):
-            caption = full.split("### Response", 1)[1].strip()
-        else:
-            caption = full[input_len:].strip()
-
-        st.info(caption or "⚠️ The LLM did not generate any text.")
-
-    except Exception as e:
-        st.error(f"Something went wrong: {e}")
-
-# Footer
-st.markdown("---")
-st.markdown("Built with ❤️ using Streamlit and Hugging Face by **shingguy1**")
+        return model_vit, tokenizer_llm, model_llm, device
+
+    model_vit, tokenizer_llm, model_llm, device = load_models()
+
+    # 6. Image uploader
+    uploaded_file = st.file_uploader("Upload a food image...", type=["jpg", "jpeg", "png"])
+
+    if uploaded_file is not None:
+        try:
+            # Display image
+            image = Image.open(uploaded_file)
+            st.image(image, caption="Uploaded Image", use_column_width=True)
+
+            # Classify with ViT
+            input_tensor = manual_transform(image).unsqueeze(0).to(device)
+            with torch.no_grad():
+                outputs = model_vit(pixel_values=input_tensor)
+            pred_idx = outputs.logits.argmax(-1).item()
+            pred_label = model_vit.config.id2label[pred_idx]
+            st.success(f"🍴 Predicted Food: **{pred_label}**")
+
+            # Build prompt
+            prompt = (
+                "### Instruction\n"
+                f"Provide a concise nutritional overview for a {pred_label}, including:\n"
+                "- Serving size (measurements & ingestion guidelines)\n"
+                "- Calories\n"
+                "- Protein, carbohydrates, and fat\n"
+                "- Main ingredients\n"
+                "- Cooking method\n"
+                "- One healthy substitution\n"
+                "### Response"
+            )
+            st.subheader("🧾 Nutrition Information")
+            st.write(f"🤖 Prompt to LLM:\n\n{prompt}")
+
+            # Tokenize & generate
+            inputs = tokenizer_llm(prompt, return_tensors="pt")
+            inputs = {k: v.to(device) for k, v in inputs.items()}
+            input_len = inputs["input_ids"].shape[1]
+
+            outputs = model_llm.generate(
+                **inputs,
+                max_length=input_len + 150,
+                temperature=0.7,
+                top_p=0.9,
+                do_sample=True,
+                no_repeat_ngram_size=2,
+                early_stopping=True,
+                pad_token_id=tokenizer_llm.eos_token_id,
+                eos_token_id=tokenizer_llm.eos_token_id
+            )
+
+            # Decode and strip prompt
+            full = tokenizer_llm.decode(outputs[0], skip_special_tokens=True).strip()
+            if "### Response" in full:
+                caption = full.split("### Response", 1)[1].strip()
+            else:
+                caption = full[input_len:].strip()
+
+            if caption:
+                st.info(caption)
+            else:
+                st.error("⚠️ The LLM did not generate any text.")
+
+        except Exception as e:
+            st.error(f"Something went wrong: {e}")
+
+if __name__ == "__main__":
+    main()