File size: 2,315 Bytes
96dd2db |
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 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 |
import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer
from azure.storage.blob import BlobServiceClient
from flask import Flask, request, jsonify
app = Flask(__name__)
# BERT model and tokenizer
model_name = "textattack/bert-base-uncased-yelp-polarity"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(model_name)
# Predict the category
def predict_category(input_text):
inputs = tokenizer(input_text, return_tensors="pt", padding=True, truncation=True)
with torch.no_grad():
logits = model(**inputs).logits
probabilities = logits.softmax(dim=1)
predicted_category = ["Documentation", "Content", "Memes"][torch.argmax(probabilities)]
return predicted_category
# Function to extract text from JSON and predict the category
def predict_category_from_json(json_data):
input_text = json_data.get('text', '')
category = predict_category(input_text)
return category
# Importing data from blob storage
def import_data_from_blob(blob_service_client, container_name, blob_name):
blob_client = blob_service_client.get_blob_client(container=container_name, blob=blob_name)
blob_data = blob_client.download_blob()
content = blob_data.readall()
return content
@app.route('/predict_category', methods=['POST'])
def predict_category_api():
try:
# Assuming JSON format with a key named 'text' that contains the text data.
json_data = request.get_json()
input_text = json_data.get('text', '')
# Predict the category
category = predict_category(input_text)
response = {'category': category}
return jsonify(response)
except Exception as e:
return jsonify({'error': str(e)})
if __name__ == '__main__':
# Azure Blob Storage connection string
connection_string = "DefaultEndpointsProtocol=https;AccountName=keywisestorage;AccountKey=uRzlCQwv/SSF6WgkEz0g83dBjnFrziSNNt8PIY5Nnt+OJic0v5xjPnO8ZMhb7SjyesYSOK79TbJ/+AStdLKiDw==;EndpointSuffix=core.windows.net"
blob_service_client = BlobServiceClient.from_connection_string(connection_string)
# Define your container and blob name
container_name = "keywisestorage"
blob_name = "pagescontainer"
app.run(host="0.0.0.0", port=5000)
|