Update app.py

app.py

@@ -1,11 +1,13 @@
 import torch
 from transformers import AutoTokenizer, AutoModelForSequenceClassification
 import gradio as gr
+import json  # Added for JSON conversion
 
 
 def run_inference(review_text: str) -> str:
     """
-    Perform inference on the given wine review text and return the predicted wine variety.
+    Perform inference on the given wine review text and return the predicted wine variety
+    using ModernBERT, an encoder-only classifier from "spawn99/modernbert-wine-classification".
 
     Args:
         review_text (str): Wine review text in the format "country [SEP] description".
@@ -19,6 +21,7 @@ def run_inference(review_text: str) -> str:
 
     # Load tokenizer and model
     tokenizer = AutoTokenizer.from_pretrained(tokenizer_id)
+    # The model used here is a ModernBERT encoder-only classifier.
     model = AutoModelForSequenceClassification.from_pretrained(model_id)
 
     # Tokenize the input text
@@ -46,26 +49,37 @@ def run_inference(review_text: str) -> str:
     return variety
 
 
-def predict_wine_variety(country: str, description: str) -> dict:
+def predict_wine_variety(country: str, description: str, output_format: str) -> str:
     """
-    Combine the provided country and description, then perform inference.
+    Combine the provided country and description, perform inference, and format the output
+    based on the selected output format.
+
     Enforces a maximum character limit of 750 on the description.
 
     Args:
         country (str): The country of wine origin.
         description (str): The wine review description.
+        output_format (str): Either "JSON" to return output as a JSON-formatted string,
+                             or "Text" for plain text output.
 
     Returns:
-        dict: Dictionary containing the predicted wine variety or an error message if the limit is exceeded.
+        str: The predicted wine variety formatted as JSON (if selected) or as plain text.
     """
-    # Validate description length
     if len(description) > 750:
-        return {"error": "Description exceeds 750 character limit. Please shorten your input."}
+        error_msg = "Description exceeds 750 character limit. Please shorten your input."
+        if output_format.lower() == "json":
+            return json.dumps({"error": error_msg}, indent=2)
+        else:
+            return error_msg
     
     # Capitalize input values and format the review text accordingly.
     review_text = f"{country.capitalize()} [SEP] {description.capitalize()}"
     predicted_variety = run_inference(review_text)
-    return {"Variety": predicted_variety}
+    
+    if output_format.lower() == "json":
+        return json.dumps({"Variety": predicted_variety}, indent=2)
+    else:
+        return predicted_variety
 
 
 if __name__ == "__main__":
@@ -73,10 +87,18 @@ if __name__ == "__main__":
         fn=predict_wine_variety,
         inputs=[
             gr.Textbox(label="Country", placeholder="Enter country of origin..."),
-            gr.Textbox(label="Description", placeholder="Enter wine review description...")
+            gr.Textbox(label="Description", placeholder="Enter wine review description..."),
+            # New radio input to choose between JSON and plain text output formats:
+            gr.Radio(choices=["JSON", "Text"], value="JSON", label="Output Format")
         ],
-        outputs=gr.JSON(label="Prediction"),
+        # Changed outputs to a Textbox so that plain text output shows naturally
+        outputs=gr.Textbox(label="Prediction"),
         title="Wine Variety Predictor",
-        description="Predict the wine variety based on country and description."
+        description=(
+            "Predict the wine variety based on the country and wine review.\n\n"
+            "This tool uses ModernBERT, an encoder-only classifier, trained on the wine reviews dataset\n"
+            "(model: spawn99/modernbert-wine-classification, dataset: spawn99/wine-reviews).\n\n"
+            "Use the Output Format selector to toggle between a JSON-formatted result and a plain text prediction."
+        )
     )
     iface.launch()