simpler and better example inference

README.md

@@ -28,52 +28,35 @@ this predicts the `ret` column of the training dataset, given the `text` column.
 
 ```py
 import json
-
-import numpy as np
-import torch
+from transformers import pipeline
 from huggingface_hub import hf_hub_download
-from transformers import AutoModelForSequenceClassification, AutoTokenizer
 
-# Define the model repository on Hugging Face Hub
 model_repo_name = "pszemraj/deberta-v3-small-sp500-edgar-10k"
+pipe = pipeline("text-classification", model=model_repo_name)
 
 # Download the regression_config.json file
 regression_config_path = hf_hub_download(
     repo_id=model_repo_name, filename="regression_config.json"
 )
-
-# Load regression configuration
 with open(regression_config_path, "r") as f:
     regression_config = json.load(f)
 
-# Load the tokenizer and model
-tokenizer = AutoTokenizer.from_pretrained(model_repo_name)
-model = AutoModelForSequenceClassification.from_pretrained(model_repo_name)
-
-
-# Function to apply inverse scaling to a prediction
 def inverse_scale(prediction, config):
+    """apply inverse scaling to a prediction"""
     min_value, max_value = config["min_value"], config["max_value"]
     return prediction * (max_value - min_value) + min_value
 
+def predict_with_pipeline(text, pipe, config, ndigits=4):
+    result = pipe(text)[0]  # Get the first (and likely only) result
+    score = result['score'] if result['label'] == 'LABEL_1' else 1 - result['score']
+    # Apply inverse scaling
+    scaled_score = inverse_scale(score, config)
+    return round(scaled_score, ndigits)
 
-# Example of using the model for inference
-def predict(text, tokenizer, model, config, ndigits=4):
-    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True)
-    with torch.no_grad():
-        outputs = model(**inputs)
-        logits = outputs.logits
-        predictions = logits.numpy()
-        # Assuming regression task, apply inverse scaling
-        scaled_predictions = [inverse_scale(pred[0], config) for pred in predictions]
-    return round(scaled_predictions[0], ndigits)
-
-
-# Example text
 text = "This is an example text for regression prediction."
 
 # Get predictions
-predictions = predict(text, tokenizer, model, regression_config)
+predictions = predict_with_pipeline(text, pipe, regression_config)
 print("Predicted Value:", predictions)
 ```