from flask import Flask, request, jsonify, render_template
from transformers import BertTokenizer, BertForSequenceClassification
import torch

app = Flask(__name__)

# Load the model and tokenizer
model = BertForSequenceClassification.from_pretrained('./model1')
tokenizer = BertTokenizer.from_pretrained('./model1')

def predict_sentiment(text):
    inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True)
    outputs = model(**inputs)
    probs = torch.nn.functional.softmax(outputs.logits, dim=1)
    pred = torch.argmax(probs).item()
    
    if pred == 0:
        sentiment = "negative"
    elif pred == 1:
        sentiment = "neutral"
    else:
        sentiment = "positive"
        
    return sentiment
    return sentiment

@app.route('/')
def home():
    return render_template('index.html')

@app.route('/predict', methods=['POST'])
def predict():
    data = request.json
    review = data['review']
    sentiment = predict_sentiment(review)
    return jsonify({'sentiment': sentiment})

if __name__ == '__main__':
    app.run(debug=True)