Update src/streamlit_app.py

src/streamlit_app.py

@@ -12,17 +12,17 @@ import torchvision.transforms as transforms
 from transformers import (
     ViTForImageClassification,
     AutoTokenizer,
-    AutoModelForCausalLM
+    T5ForConditionalGeneration
 )
 
 def main():
-    # Environment & cache
+    # 1. 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
 
-    # Image transform for ViT
+    # 2. Image transform for ViT
     manual_transform = transforms.Compose([
         transforms.Resize(256),
         transforms.CenterCrop(224),
@@ -33,102 +33,94 @@ def main():
         transforms.ConvertImageDtype(torch.float32)
     ])
 
-    # Sidebar info
+    # 3. Sidebar info
     st.sidebar.header("Models Used")
     st.sidebar.markdown("""
     - 🖼️ **Image Classifier**: `shingguy1/fine_tuned_vit`  
-    - 💬 **Text Generator**: `tiiuae/falcon-7b-instruct`
+    - 💬 **Text Generator**: `google/flan-t5-small`
     """)
 
-    # Load models (cached)
+    # 4. Load models (cached)
     @st.cache_resource
     def load_models():
-        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        device = torch.device("cpu")  # CPU-only environment
 
-        # ViT classifier → GPU/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 → 8-bit quant on GPU
+        # FLAN-T5 Small for generation
         tokenizer_llm = AutoTokenizer.from_pretrained(
-            "tiiuae/falcon-7b-instruct",
+            "google/flan-t5-small",
             cache_dir=cache_dir,
             use_auth_token=hf_token
         )
-        model_llm = AutoModelForCausalLM.from_pretrained(
-            "tiiuae/falcon-7b-instruct",
+        model_llm = T5ForConditionalGeneration.from_pretrained(
+            "google/flan-t5-small",
             cache_dir=cache_dir,
-            load_in_8bit=True,
-            device_map="auto",
-            torch_dtype=torch.float16,
             use_auth_token=hf_token
-        )
+        ).to(device)
 
         return model_vit, tokenizer_llm, model_llm, device
 
     model_vit, tokenizer_llm, model_llm, device = load_models()
 
-    # Image uploader
+    # 5. 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
-            inputs_v = manual_transform(image).unsqueeze(0).to(device)
+            # Classify with ViT
+            inputs_vit = manual_transform(image).unsqueeze(0).to(device)
             with torch.no_grad():
-                out = model_vit(pixel_values=inputs_v)
-            idx = out.logits.argmax(-1).item()
-            pred_label = model_vit.config.id2label[idx]
+                vit_outputs = model_vit(pixel_values=inputs_vit)
+            pred_idx = vit_outputs.logits.argmax(-1).item()
+            pred_label = model_vit.config.id2label[pred_idx]
             st.success(f"🍴 Predicted Food: **{pred_label}**")
 
-            # Unified instruction prompt
+            # Build FLAN-T5 prompt
             prompt = (
-                "### Instruction\n"
-                f"Provide a concise nutritional overview for a {pred_label}, including:\n"
-                "- Serving size (measurements & ingestion guidelines)\n"
+                "Provide a concise nutritional overview for a taco, including:\n"
+                "- Serving size (with measurements & ingestion guidelines)\n"
                 "- Calories\n"
                 "- Protein, carbohydrates, and fat\n"
                 "- Main ingredients\n"
                 "- Cooking method\n"
                 "- One healthy substitution\n"
-                "### Response"
+                "Answer only the overview."
             )
             st.subheader("🧾 Nutrition Information")
-            st.write(f"🤖 Prompt to LLM:\n\n{prompt}")
+            st.write(f"🤖 Prompt:\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()}
-            inp_len = inputs["input_ids"].shape[1]
-
-            out_ids = model_llm.generate(
-                **inputs,
-                max_length=inp_len + 150,
+            inputs = tokenizer_llm(
+                prompt,
+                return_tensors="pt",
+                padding="longest",
+                truncation=True,
+            ).to(device)
+
+            outputs = model_llm.generate(
+                input_ids=inputs.input_ids,
+                attention_mask=inputs.attention_mask,
+                max_new_tokens=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,
+                pad_token_id=tokenizer_llm.pad_token_id,
                 eos_token_id=tokenizer_llm.eos_token_id
-            )[0]
-
-            # Decode & strip prompt
-            decoded = tokenizer_llm.decode(out_ids, skip_special_tokens=True).strip()
-            if "### Response" in decoded:
-                caption = decoded.split("### Response", 1)[1].strip()
-            else:
-                caption = decoded[inp_len:].strip()
-
-            if caption:
-                st.info(caption)
-            else:
-                st.error("⚠️ The LLM did not generate any text.")
+            )
+
+            summary = tokenizer_llm.decode(outputs[0], skip_special_tokens=True).strip()
+            st.info(summary or "⚠️ The model did not generate any text.")
 
         except Exception as e:
             st.error(f"Something went wrong: {e}")