File size: 896 Bytes
82e0279 9bc3677 82e0279 db6e2f8 53f0da3 db6e2f8 82e0279 db6e2f8 82e0279 db6e2f8 674a5ba db6e2f8 82e0279 db6e2f8 674a5ba db6e2f8 674a5ba db6e2f8 82e0279 674a5ba db6e2f8 82e0279 674a5ba 82e0279 db6e2f8 | 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 | 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
|