Update app.py

app.py

@@ -2,7 +2,7 @@ import gradio as gr
 from transformers import pipeline
 import torch
 
-# List of your models (using repo IDs for easy loading)
+# List of models (unchanged)
 MODELS = {
     "econbert": "climatebert/econbert",
     "controversy-classification": "climatebert/ClimateControversyBERT_classification",
@@ -13,46 +13,47 @@ MODELS = {
     "climate-detector": "climatebert/distilroberta-base-climate-detector",
     "climate-commitment": "climatebert/distilroberta-base-climate-commitment",
     "climate-tcfd": "climatebert/distilroberta-base-climate-tcfd",
-    "climate-s": "climatebert/distilroberta-base-climate-s",  # Adjust label if truncated
+    "climate-s": "climatebert/distilroberta-base-climate-s",
     "climate-specificity": "climatebert/distilroberta-base-climate-specificity",
     "climate-sentiment": "climatebert/distilroberta-base-climate-sentiment",
     "environmental-claims": "climatebert/environmental-claims",
-    "climate-f": "climatebert/distilroberta-base-climate-f",  # Adjust label if truncated
-    "climate-d-s": "climatebert/distilroberta-base-climate-d-s",  # Adjust label if truncated
-    "climate-d": "climatebert/distilroberta-base-climate-d"   # Adjust label if truncated
+    "climate-f": "climatebert/distilroberta-base-climate-f",
+    "climate-d-s": "climatebert/distilroberta-base-climate-d-s",
+    "climate-d": "climatebert/distilroberta-base-climate-d"
 }
 
-# Cache for loaded pipelines (lazy loading)
+# Cache for loaded pipelines
 pipelines = {}
 
 def load_model(model_key):
-    """Load pipeline for the selected model (text-classification assumed; adjust task if needed)."""
+    """Load pipeline for the selected model with truncation enabled."""
     if model_key not in pipelines:
         repo_id = MODELS[model_key]
-        device = 0 if torch.cuda.is_available() else -1  # GPU if available
+        device = 0 if torch.cuda.is_available() else -1
         pipelines[model_key] = pipeline(
             "text-classification",
             model=repo_id,
             device=device,
-            torch_dtype=torch.float16 if device == 0 else None  # Half-precision on GPU for memory savings
+            torch_dtype=torch.float16 if device == 0 else None,
+            truncation=True,  # Truncate input to max_length
+            max_length=512   # Set to 512 to match BERT/DistilRoBERTa limit
         )
     return pipelines[model_key]
 
 def predict(model_key, text):
-    """Run inference on selected model."""
+    """Run inference on selected model with truncation."""
     if not text.strip():
         return "Please enter some text."
     
     try:
         model = load_model(model_key)
         results = model(text)
-        # Format output (e.g., for classification: label + score)
         formatted = "\n".join([f"{r['label']}: {r['score']:.2f}" for r in results])
-        return f"Predictions for '{text}':\n{formatted}"
+        return f"Predictions for '{text[:50]}...':\n{formatted}"  # Truncate display for readability
     except Exception as e:
-        return f"Error: {str(e)} (Check model card for task/input format)."
+        return f"Error: {str(e)} (Check input length or model card for details)."
 
-# Gradio interface
+# Gradio interface (unchanged)
 with gr.Blocks(title="ClimateBERT Multi-Model Demo") as demo:
     gr.Markdown("# ClimateBERT Models Demo\nSelect a model and input text for climate-related analysis (e.g., sentiment, classification).")