File size: 2,340 Bytes
d54d478 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 |
"""
Example: Using the Energy Document Classifier
This script demonstrates how to use the model for classifying documents.
"""
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch
def main():
# Load model and tokenizer
model_name = "EnergyAI/Llama-3.1-8B-Energy-Classifier" # Change to your model
print("Loading model...")
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(
model_name,
device_map="auto",
torch_dtype=torch.bfloat16,
)
model.eval()
print("Model loaded!\n")
# Example texts
texts = [
"Solar panel installations have increased by 40% this year, driven by government incentives and falling prices.",
"The software development team completed the sprint planning session and assigned tasks for the next iteration.",
"OPEC announced a production cut of 2 million barrels per day, causing oil prices to surge on global markets.",
"The training program for new employees will begin next Monday and continue for three weeks.",
]
# Classify each text
label_map = {0: "non_energy", 1: "energy"}
for i, text in enumerate(texts, 1):
print(f"\n{'='*70}")
print(f"Example {i}:")
print(f"Text: {text}")
print(f"{'-'*70}")
# Tokenize and prepare input
inputs = tokenizer(
text,
return_tensors="pt",
truncation=True,
max_length=512,
padding=True,
)
inputs = {k: v.to(model.device) for k, v in inputs.items()}
# Get prediction
with torch.no_grad():
outputs = model(**inputs)
probs = torch.nn.functional.softmax(outputs.logits, dim=-1)
predicted_class = torch.argmax(probs, dim=-1).item()
confidence = probs[0][predicted_class].item()
# Print results
print(f"Prediction: {label_map[predicted_class].upper()}")
print(f"Confidence: {confidence:.4f}")
print(f"Probabilities:")
print(f" - Non-Energy: {probs[0][0].item():.4f}")
print(f" - Energy: {probs[0][1].item():.4f}")
print(f"\n{'='*70}\n")
if __name__ == "__main__":
main()
|