Update src/streamlit_app.py

src/streamlit_app.py

@@ -53,7 +53,7 @@ def main():
     st.sidebar.header("Models Used")
     st.sidebar.markdown("""
     - 🖼️ **Image Classifier**: `shingguy1/fine_tuned_vit`
-    - 💬 **Paraphraser**: `google/flan-t5-small` (beam search mode)
+    - 💬 **Paraphraser**: `google/flan-t5-small` (sampling mode)
     """)
 
     transform = transforms.Compose([
@@ -68,9 +68,21 @@ def main():
     @st.cache_resource
     def load_models():
         device = torch.device("cpu")
-        vit = ViTForImageClassification.from_pretrained("shingguy1/fine_tuned_vit", cache_dir=cache_dir, use_auth_token=hf_token).to(device)
-        tok = AutoTokenizer.from_pretrained("google/flan-t5-small", cache_dir=cache_dir, use_auth_token=hf_token)
-        t5 = T5ForConditionalGeneration.from_pretrained("google/flan-t5-small", cache_dir=cache_dir, use_auth_token=hf_token).to(device)
+        vit = ViTForImageClassification.from_pretrained(
+            "shingguy1/fine_tuned_vit",
+            cache_dir=cache_dir,
+            use_auth_token=hf_token
+        ).to(device)
+        tok = AutoTokenizer.from_pretrained(
+            "google/flan-t5-small",
+            cache_dir=cache_dir,
+            use_auth_token=hf_token
+        )
+        t5 = T5ForConditionalGeneration.from_pretrained(
+            "google/flan-t5-small",
+            cache_dir=cache_dir,
+            use_auth_token=hf_token
+        ).to(device)
         return vit, tok, t5, device
 
     model_vit, tokenizer_t5, model_t5, device = load_models()
@@ -97,18 +109,24 @@ def main():
                 f"Try {data['substitute']} as a healthier swap."
             )
             prompt = (
-                f"Rewrite the following nutritional information in a clear and friendly tone. "
-                f"Do not add or change any facts. Keep the sentence structure simple:\n\n"
+                f"Paraphrase the following nutritional facts in a friendly, conversational tone. "
+                f"Use varied sentence structures and synonyms, and feel free to generalize numeric details "
+                f"(e.g., ‘around 250 kcal’). Don’t add any new facts.\n\n"
                 f"{base_description}"
             )
         else:
-            prompt = f"Give the typical calories, macros, and nutrition facts for {label}. Provide realistic values even if estimated."
+            prompt = (
+                f"Provide an approximate nutrition summary for {label}, including calories, "
+                f"macronutrients, and a brief description."
+            )
 
         inputs = tokenizer_t5(prompt, return_tensors="pt", truncation=True).to(device)
         output_ids = model_t5.generate(
             inputs["input_ids"],
             max_new_tokens=100,
-            num_beams=4,
+            do_sample=True,
+            top_p=0.9,
+            temperature=0.7,
             early_stopping=True
         )
         response = tokenizer_t5.decode(output_ids[0], skip_special_tokens=True)