from flask import Flask, request
from transformers import AutoModelForSequenceClassification, AutoTokenizer, TextClassificationPipeline
import torch

app = Flask(__name__)

tokenizer = AutoTokenizer.from_pretrained('alexray/twitter-distilbert')

model = AutoModelForSequenceClassification.from_pretrained('alexray/twitter-distilbert', num_labels=2)

# Create a text classification pipeline
classifier = TextClassificationPipeline(model=model, tokenizer=tokenizer, device=-1)  # device=-1 means using CPU

@app.route('/')
def predict() -> str:
    args = request.args
    data = args.get("data")

    if data is None:
        return 'data is none'

    # Make prediction using the pipeline
    result = classifier(data)

    # Get predicted class (assuming it's a binary classification)
    predicted_class = result[0]['label']

    http_response = str(predicted_class)

    return http_response