import joblib
import numpy as np
from sentence_transformers import SentenceTransformer

# Load the models
model1 = joblib.load('model1.joblib')
model2 = joblib.load('model2.joblib')

# Load the embedder
embedder = SentenceTransformer('BAAI/bge-large-en-v1.5')

def predict_sentiment(text):
    # Generate embedding
    embedding = embedder.encode([text])
    
    # Make predictions
    pred1 = model1.predict(embedding)[0]
    pred2 = model2.predict(embedding)[0]
    
    # Average and round
    final_prediction = np.round((pred1 + pred2) / 2).astype(int)
    
    return final_prediction, pred1, pred2

# Example usage
if __name__ == "__main__":
    test_text = "I really enjoyed this product!"
    final_score, score1, score2 = predict_sentiment(test_text)
    print(f"Text: {test_text}")
    print(f"Final sentiment score: {final_score}")
    print(f"Model 1 score: {score1}")
    print(f"Model 2 score: {score2}")