Nechba's picture
Update app.py
68d753f verified
raw
history blame
2.33 kB
from flask import Flask, request, jsonify
from transformers import pipeline
from transformers import AutoTokenizer, AutoModelForTokenClassification
import whisper
import os
import tempfile
app = Flask(__name__)
# Initialize Whisper model
whisper_model = whisper.load_model("small") # Renamed variable
# Initialize Emotion Classifier
classifier = pipeline("text-classification", model="j-hartmann/emotion-english-distilroberta-base", return_all_scores=True)
# Initialize NER pipeline
ner_tokenizer = AutoTokenizer.from_pretrained("dslim/bert-base-NER")
ner_model = AutoModelForTokenClassification.from_pretrained("dslim/bert-base-NER") # Renamed variable
ner_pipeline = pipeline("ner", model=ner_model, tokenizer=ner_tokenizer) # Renamed variable
@app.route('/transcribe', methods=['POST'])
def transcribe():
try:
# Read the raw audio bytes
audio_bytes = request.data # Get raw bytes from request
if not audio_bytes:
return jsonify({"error": "No audio data provided"}), 400
# Convert bytes to a file-like object
audio_file = io.BytesIO(audio_bytes)
# Transcribe the audio
result = whisper_model.transcribe(audio_file)
return jsonify({"text": result["text"]})
except Exception as e:
print("Error:", str(e)) # Log the error
return jsonify({"error": "Internal Server Error", "details": str(e)}), 500
@app.route('/classify', methods=['POST'])
def classify():
try:
data = request.get_json()
if 'text' not in data:
return jsonify({"error": "Missing 'text' field"}), 400
text = data['text']
result = classifier(text)
return jsonify(result)
except Exception as e:
return jsonify({"error": str(e)}), 500
@app.route('/ner', methods=['POST'])
def ner_endpoint():
try:
data = request.get_json()
text = data.get("text", "")
# Use the renamed ner_pipeline
ner_results = ner_pipeline(text)
words_and_entities = [
{"word": result['word'], "entity": result['entity']}
for result in ner_results
]
return jsonify({"entities": words_and_entities})
except Exception as e:
return jsonify({"error": str(e)}), 500