karthikmn's picture
Create model.py
108667c verified
raw
history blame
1.34 kB
import torch
from transformers import BertTokenizer, BertModel, T5ForConditionalGeneration
from transformers import AutoModelForSequenceClassification
from utilities import preprocess_features, generate_recommendation
MODEL_PATH = "your-hf-username/deal-score-model" # Replace with your model path
SUMMARIZER_PATH = "t5-small" # Can use distilgpt2 or other summarizers
def load_model():
model = AutoModelForSequenceClassification.from_pretrained(MODEL_PATH)
tokenizer = BertTokenizer.from_pretrained(MODEL_PATH)
summarizer = T5ForConditionalGeneration.from_pretrained(SUMMARIZER_PATH)
return model, tokenizer, summarizer
def predict_deal_score(input_data, model, tokenizer, summarizer):
features = preprocess_features(input_data)
inputs = tokenizer(features["text"], return_tensors="pt")
with torch.no_grad():
outputs = model(**inputs)
score = round(outputs.logits[0][0].item() * 100)
confidence = torch.softmax(outputs.logits[0], dim=-1).max().item()
risk = (
"Low" if score > 75 else
"Medium" if score > 50 else
"High"
)
recommendation = generate_recommendation(summarizer, input_data)
return {
"score": score,
"confidence": round(confidence, 2),
"risk": risk,
"recommendation": recommendation
}