new model

app.py

@@ -1,36 +1,24 @@
 """
 =========================================================
  app.py — Green Greta (Gradio + TF/Keras 3 + LangChain v0.2)
- - Image model: load TF SavedModel via keras.layers.TFSMLayer (Keras 3 safe)
- - LLM: local transformers pipeline (no HF API token required)
- - LangChain v0.2 imports (text_splitters/core/community)
- - Robust JSON parsing for schema-shaped output
- - EfficientNet input size fix (224x224)
- - Gradio binds to 0.0.0.0:7860 (Docker-friendly)
 =========================================================
 """
 
-# =========================
-# Imports (grouped together)
-# =========================
+# ========== Imports ==========
 import os
 import json
 import shutil
 
-# UI / web
 import gradio as gr
-
-# TensorFlow / Keras / image
 import tensorflow as tf
 from tensorflow import keras
 from PIL import Image
 
-# Networking / retry
 import tenacity
 from fake_useragent import UserAgent
 
 # LangChain v0.2 family
-from langchain_text_splitters import RecursiveCharacterTextSplitter
+from langchain_text_splitters import RecursiveCharacterTextSplitter, TokenTextSplitter
 from langchain_core.prompts import ChatPromptTemplate
 from langchain_core.output_parsers import PydanticOutputParser
 from langchain_community.document_loaders import WebBaseLoader
@@ -38,67 +26,64 @@ from langchain_community.vectorstores import Chroma
 from langchain.chains import ConversationalRetrievalChain
 from langchain.memory import ConversationBufferMemory
 
-# Embeddings (community version works; you can switch to langchain-huggingface later)
-from langchain_community.embeddings import HuggingFaceEmbeddings
-# If you prefer to silence deprecation warnings in the future:
-# from langchain_huggingface import HuggingFaceEmbeddings  # pip install -U langchain-huggingface
+# Embeddings (use community; switch to langchain-huggingface later if desired)
+try:
+    from langchain_huggingface import HuggingFaceEmbeddings  # pip install -U langchain-huggingface
+except ImportError:
+    from langchain_community.embeddings import HuggingFaceEmbeddings
+
+# Context compression (keeps inputs ≤ model limit)
+from langchain.retrievers import ContextualCompressionRetriever
+from langchain.retrievers.document_compressors import DocumentCompressorPipeline
 
-# Pydantic for schema in prompt
+# Pydantic schema
 from pydantic.v1 import BaseModel, Field
 
-# Hugging Face Hub helper for SavedModel
+# HF Hub for SavedModel download
 from huggingface_hub import snapshot_download
 
-# Local theming + URLs list
+# Transformers local pipeline (no token needed)
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
+from langchain_community.llms import HuggingFacePipeline
+
+# Local theme + URLs
 import theme
 from url_list import URLS
 
 
-# =========================
-# Theme instance
-# =========================
+# ========== Theme ==========
 theme = theme.Theme()
 
 
 # =========================================================
-# 1) IMAGE CLASSIFICATION MODEL SETUP (Keras 3-compatible)
+# 1) IMAGE CLASSIFICATION (Keras 3-compatible SavedModel)
 # =========================================================
-# The HF repo is a TensorFlow SavedModel; with Keras 3 we must use TFSMLayer.
 MODEL_REPO = "rocioadlc/efficientnetB0_trash"
-MODEL_SERVING_SIGNATURE = "serving_default"  # adjust if your repo uses another signature
+MODEL_SERVING_SIGNATURE = "serving_default"  # adjust if your model uses a different signature
 
-# Download SavedModel locally
+# Download the SavedModel once and wrap with Keras TFSMLayer
 model_dir = snapshot_download(MODEL_REPO)
-
-# Wrap SavedModel as a Keras layer
 model1 = keras.layers.TFSMLayer(model_dir, call_endpoint=MODEL_SERVING_SIGNATURE)
 
-# Class labels
 class_labels = ["cardboard", "glass", "metal", "paper", "plastic", "trash"]
 
-
 def predict_image(input_image: Image.Image):
-    """
-    Resize the user-uploaded image and preprocess it for EfficientNetB0.
-    Works with a TFSMLayer (SavedModel) that may return a dict of tensors.
-    """
-    img = input_image.convert("RGB").resize((224, 224))  # EfficientNetB0 expects 224x224
+    """Preprocess to 224x224 EfficientNet input and run inference."""
+    img = input_image.convert("RGB").resize((224, 224))
     x = tf.keras.preprocessing.image.img_to_array(img)
     x = tf.keras.applications.efficientnet.preprocess_input(x)
-    x = tf.expand_dims(x, 0)  # [1, 224, 224, 3]
+    x = tf.expand_dims(x, 0)  # batch
 
     outputs = model1(x)
     if isinstance(outputs, dict) and outputs:
-        key = next(iter(outputs))
-        preds = outputs[key]
+        preds = outputs[next(iter(outputs))]
     else:
         preds = outputs
 
-    preds_np = preds.numpy() if hasattr(preds, "numpy") else preds
-    probs = preds_np[0].tolist()
+    arr = preds.numpy() if hasattr(preds, "numpy") else preds
+    probs = arr[0].tolist()
     return {label: float(probs[i]) for i, label in enumerate(class_labels)}
 
-
 image_gradio_app = gr.Interface(
     fn=predict_image,
     inputs=gr.Image(label="Image", sources=["upload", "webcam"], type="pil"),
@@ -114,13 +99,11 @@ image_gradio_app = gr.Interface(
 user_agent = UserAgent().random
 header_template = {"User-Agent": user_agent}
 
-
 @tenacity.retry(wait=tenacity.wait_fixed(3), stop=tenacity.stop_after_attempt(3), reraise=True)
 def load_url(url: str):
     loader = WebBaseLoader(web_paths=[url], header_template=header_template)
     return loader.load()
 
-
 def safe_load_all_urls(urls):
     all_docs = []
     for link in urls:
@@ -131,21 +114,20 @@ def safe_load_all_urls(urls):
             print(f"Skipping URL due to error: {link}\nError: {e}\n")
     return all_docs
 
-
 all_loaded_docs = safe_load_all_urls(URLS)
 
+# Smaller base chunks to help keep prompts short
 text_splitter = RecursiveCharacterTextSplitter(
-    chunk_size=1024,
-    chunk_overlap=150,
+    chunk_size=700,
+    chunk_overlap=80,
     length_function=len,
 )
-
 docs = text_splitter.split_documents(all_loaded_docs)
 
 # Embeddings
 embeddings = HuggingFaceEmbeddings(model_name="thenlper/gte-small")
 
-# Vector store (Chroma)
+# Vector store
 persist_directory = "docs/chroma/"
 shutil.rmtree(persist_directory, ignore_errors=True)
 
@@ -155,7 +137,17 @@ vectordb = Chroma.from_documents(
     persist_directory=persist_directory,
 )
 
-retriever = vectordb.as_retriever(search_kwargs={"k": 3}, search_type="mmr")
+# Base retriever
+base_retriever = vectordb.as_retriever(search_kwargs={"k": 2}, search_type="mmr")
+
+# Hard-cap tokens in retrieved docs (~200 tokens per slice)
+token_splitter = TokenTextSplitter(chunk_size=200, chunk_overlap=30)
+compressor = DocumentCompressorPipeline(transformers=[token_splitter])
+
+compression_retriever = ContextualCompressionRetriever(
+    base_retriever=base_retriever,
+    base_compressor=compressor,
+)
 
 
 # ======================================
@@ -165,23 +157,15 @@ class FinalAnswer(BaseModel):
     question: str = Field(description="User question")
     answer: str = Field(description="Direct answer")
 
-
 parser = PydanticOutputParser(pydantic_object=FinalAnswer)
 
 SYSTEM_TEMPLATE = (
-    """
-Your name is Greta and you are a recycling chatbot with the objective to answer questions from user in English or Spanish /
-Has sido diseñado y creado por el Grupo 1 del Máster en Data Science & Big Data de la promoción 2023/2024 de la Universidad Complutense de Madrid. Este grupo está formado por Rocío, María Guillermo, Alejandra, Paloma y Álvaro /
-Use the following pieces of context to answer the question /
-If the question is English answer in English /
-If the question is Spanish answer in Spanish /
-Do not mention the word context when you answer a question /
-Answer the question fully and provide as much relevant detail as possible. Do not cut your response short /
-Context: {context}
-User: {question}
-{format_instructions}
-    """
-).strip()
+    "You are Greta, a bilingual (EN/ES) recycling assistant. "
+    "Answer fully using the snippets below. Do not mention 'context'.\n\n"
+    "Context:\n{context}\n\n"
+    "User: {question}\n"
+    "{format_instructions}"
+)
 
 qa_prompt = ChatPromptTemplate.from_template(
     SYSTEM_TEMPLATE,
@@ -192,12 +176,7 @@ qa_prompt = ChatPromptTemplate.from_template(
 # =============================
 # 4) LLM (token-free local model)
 # =============================
-# Avoids HF Endpoint auth + deprecated .post path. Good defaults for CPU.
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
-from langchain_community.llms import HuggingFacePipeline
-
 LOCAL_MODEL_ID = os.environ.get("LOCAL_LLM", "google/flan-t5-base")
-
 tok = AutoTokenizer.from_pretrained(LOCAL_MODEL_ID)
 mdl = AutoModelForSeq2SeqLM.from_pretrained(LOCAL_MODEL_ID)
 
@@ -206,9 +185,8 @@ gen = pipeline(
     model=mdl,
     tokenizer=tok,
     max_new_tokens=512,
-    do_sample=False,
+    do_sample=False,  # deterministic; better for JSON adherence
 )
-
 llm = HuggingFacePipeline(pipeline=gen)
 
 
@@ -222,7 +200,7 @@ memory = ConversationBufferMemory(
 
 qa_chain = ConversationalRetrievalChain.from_llm(
     llm=llm,
-    retriever=retriever,
+    retriever=compression_retriever,  # <= compressed retriever to avoid 512-token overflows
     memory=memory,
     verbose=True,
     combine_docs_chain_kwargs={"prompt": qa_prompt},
@@ -231,34 +209,26 @@ qa_chain = ConversationalRetrievalChain.from_llm(
     output_key="output",
 )
 
+def _safe_json_extract(raw: str, question: str) -> dict:
+    """Try strict JSON; otherwise extract first {...} block; fallback to plain text."""
+    raw = (raw or "").strip()
+    try:
+        return json.loads(raw)
+    except json.JSONDecodeError:
+        start = raw.find("{")
+        end = raw.rfind("}")
+        if start != -1 and end != -1 and end > start:
+            try:
+                return json.loads(raw[start : end + 1])
+            except json.JSONDecodeError:
+                pass
+    return {"question": question, "answer": raw or "No answer produced."}
 
 def chat_interface(question, history):
-    """
-    Run the QA chain and return the 'answer' field from a JSON payload.
-    Falls back safely if the LLM returns non-JSON text.
-    """
     try:
         result = qa_chain.invoke({"question": question})
-        raw = result.get("output", "").strip()
-
-        # Strict JSON first
-        try:
-            payload = json.loads(raw)
-        except json.JSONDecodeError:
-            # Try extracting first {...} block
-            start = raw.find("{")
-            end = raw.rfind("}")
-            if start != -1 and end != -1 and end > start:
-                try:
-                    payload = json.loads(raw[start : end + 1])
-                except json.JSONDecodeError:
-                    payload = {"question": question, "answer": raw}
-            else:
-                payload = {"question": question, "answer": raw}
-
-        # Return the schema field
-        return payload.get("answer", raw)
-
+        payload = _safe_json_extract(result.get("output", ""), question)
+        return payload.get("answer", "")
     except Exception as e:
         return (
             "Lo siento, tuve un problema procesando tu pregunta. "
@@ -266,7 +236,6 @@ def chat_interface(question, history):
             f"Detalle técnico: {e}"
         )
 
-
 chatbot_gradio_app = gr.ChatInterface(
     fn=chat_interface,
     title="<span style='color: rgb(243, 239, 224);'>Green Greta</span>",
@@ -284,7 +253,7 @@ banner_tab_content = """
     <p style="font-size: 16px; color: #4e6339; text-align: justify;">Nuestra plataforma combina la potencia de la inteligencia artificial con la comodidad de un chatbot para brindarte respuestas rápidas y precisas sobre qué objetos son reciclables y cómo hacerlo de la manera más eficiente.</p>
     <p style="font-size: 16px; text-align:center;"><strong><span style="color: #4e6339;">¿Cómo usarlo?</span></strong></p>
     <ul style="list-style-type: disc; text-align: justify; margin-top: 20px; padding-left: 20px;">
-        <li style="font-size: 16px; color: #4e6339;"><strong><span style="color: #4e6339;">Green Greta Image Classification:</span></strong> Ve a la pestaña Greta Image Classification y simplemente carga una foto del objeto que quieras reciclar, y nuestro modelo de identificará de qué se trata🕵️‍♂️ para que puedas desecharlo adecuadamente.</li>
+        <li style="font-size: 16px; color: #4e6339;"><strong><span style="color: #4e6339;">Green Greta Image Classification:</span></strong> Ve a la pestaña Greta Image Classification y simplemente carga una foto del objeto que quieras reciclar, y nuestro modelo identificará de qué se trata🕵️‍♂️ para que puedas desecharlo adecuadamente.</li>
         <li style="font-size: 16px; color: #4e6339;"><strong><span style="color: #4e6339;">Green Greta Chat:</span></strong> ¿Tienes preguntas sobre reciclaje, materiales específicos o prácticas sostenibles? ¡Pregunta a nuestro chatbot en la pestaña Green Greta Chat!📝 Está aquí para responder todas tus preguntas y ayudarte a tomar decisiones más informadas sobre tu reciclaje.</li>
     </ul>
     <h1 style="font-size: 24px; color: #4e6339; margin-top: 20px;">Welcome to our image classifier and chatbot for smarter recycling!♻️</h1>
@@ -297,7 +266,6 @@ banner_tab_content = """
     </ul>
 </div>
 """
-
 banner_tab = gr.Markdown(banner_tab_content)
 
 
@@ -310,7 +278,6 @@ app = gr.TabbedInterface(
     theme=theme,
 )
 
-# Concurrency queue + launch (Docker-friendly binding)
 app.queue()
 app.launch(
     server_name="0.0.0.0",