Upload app.py

app.py

@@ -0,0 +1,83 @@
+import torch
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
+import gradio as gr
+
+
+def run_inference(review_text: str) -> str:
+    """
+    Perform inference on the given wine review text and return the predicted wine variety.
+
+    Args:
+        review_text (str): Wine review text in the format "country [SEP] description".
+
+    Returns:
+        str: The predicted wine variety using the model's id2label mapping if available.
+    """
+    # Define model and tokenizer identifiers
+    model_id = "spawn99/modernbert-wine-classification"
+    tokenizer_id = "answerdotai/ModernBERT-base"
+
+    # Load tokenizer and model
+    tokenizer = AutoTokenizer.from_pretrained(tokenizer_id)
+    model = AutoModelForSequenceClassification.from_pretrained(model_id)
+
+    # Tokenize the input text
+    inputs = tokenizer(
+        review_text,
+        return_tensors="pt",
+        padding="max_length",
+        truncation=True,
+        max_length=256
+    )
+
+    model.eval()
+    with torch.no_grad():
+        outputs = model(**inputs)
+    logits = outputs.logits
+
+    # Determine prediction and map to label if available
+    pred = torch.argmax(logits, dim=-1).item()
+    variety = (
+        model.config.id2label.get(pred, str(pred))
+        if hasattr(model.config, "id2label") and model.config.id2label
+        else str(pred)
+    )
+    
+    return variety
+
+
+def predict_wine_variety(country: str, description: str) -> dict:
+    """
+    Combine the provided country and description, then perform inference.
+    Enforces a maximum character limit of 750 on the description.
+
+    Args:
+        country (str): The country of wine origin.
+        description (str): The wine review description.
+
+    Returns:
+        dict: Dictionary containing the predicted wine variety or an error message if the limit is exceeded.
+    """
+    # Validate description length
+    if len(description) > 750:
+        return {"error": "Description exceeds 750 character limit. Please shorten your input."}
+    
+    # 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 __name__ == "__main__":
+    iface = gr.Interface(
+        fn=predict_wine_variety,
+        inputs=[
+            gr.Textbox(label="Country", placeholder="Enter country of origin..."),
+            gr.Textbox(label="Description", placeholder="Enter wine review description...")
+        ],
+        outputs=gr.JSON(label="Prediction"),
+        title="Wine Variety Predictor",
+        description="Predict the wine variety based on country and description.",
+        flagging="never"
+    )
+    iface.launch()