Nechba commited on
Commit
fbe647d
·
verified ·
1 Parent(s): 05c5f59

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +15 -20
app.py CHANGED
@@ -1,47 +1,46 @@
1
  from flask import Flask, request, jsonify
2
  from transformers import pipeline
3
  from transformers import AutoTokenizer, AutoModelForTokenClassification
4
-
5
- # Initialize the tokenizer and model
6
  import whisper
7
  import os
8
 
9
  app = Flask(__name__)
10
 
 
 
11
 
 
 
12
 
13
- # Load the model once at startup (better performance for multiple requests)
14
- model = whisper.load_model("small")
 
 
15
 
16
  def allowed_file(filename):
17
  return '.' in filename and filename.rsplit('.', 1)[1].lower() in {'wav', 'mp3', 'ogg', 'flac', 'm4a'}
18
 
19
  @app.route('/transcribe', methods=['POST'])
20
  def transcribe_audio():
21
- # Check if a file was uploaded
22
  if 'file' not in request.files:
23
  return jsonify({'error': 'No file uploaded'}), 400
24
 
25
  file = request.files['file']
26
 
27
- # Check if the file is empty
28
  if file.filename == '':
29
  return jsonify({'error': 'No selected file'}), 400
30
 
31
- # Check allowed file types
32
  if not allowed_file(file.filename):
33
  return jsonify({'error': 'Unsupported file type'}), 400
34
 
35
  try:
36
- # Save the temporary file
37
  temp_path = "temp_audio"
38
  file.save(temp_path)
39
 
40
- # Transcribe the audio
41
- result = model.transcribe(temp_path)
42
  transcription = result["text"]
43
 
44
- # Clean up the temporary file
45
  if os.path.exists(temp_path):
46
  os.remove(temp_path)
47
 
@@ -49,7 +48,6 @@ def transcribe_audio():
49
 
50
  except Exception as e:
51
  return jsonify({'error': str(e)}), 500
52
- classifier = pipeline("text-classification", model="j-hartmann/emotion-english-distilroberta-base", return_all_scores=True)
53
 
54
  @app.route('/classify', methods=['POST'])
55
  def classify():
@@ -65,26 +63,23 @@ def classify():
65
  except Exception as e:
66
  return jsonify({"error": str(e)}), 500
67
 
68
- tokenizer = AutoTokenizer.from_pretrained("dslim/bert-base-NER")
69
- model = AutoModelForTokenClassification.from_pretrained("dslim/bert-base-NER")
70
- nlp = pipeline("ner", model=model, tokenizer=tokenizer)
71
  @app.route('/ner', methods=['POST'])
72
  def ner_endpoint():
73
  try:
74
- # Get text from request
75
  data = request.get_json()
76
  text = data.get("text", "")
77
 
78
- # Perform NER
79
- ner_results = nlp(text)
80
 
81
- # Extract words and their corresponding entities
82
  words_and_entities = [
83
  {"word": result['word'], "entity": result['entity']}
84
  for result in ner_results
85
  ]
86
 
87
- # Return JSON response with the words and their entities
88
  return jsonify({"entities": words_and_entities})
89
  except Exception as e:
90
  return jsonify({"error": str(e)}), 500
 
 
 
 
1
  from flask import Flask, request, jsonify
2
  from transformers import pipeline
3
  from transformers import AutoTokenizer, AutoModelForTokenClassification
 
 
4
  import whisper
5
  import os
6
 
7
  app = Flask(__name__)
8
 
9
+ # Initialize Whisper model
10
+ whisper_model = whisper.load_model("small") # Renamed variable
11
 
12
+ # Initialize Emotion Classifier
13
+ classifier = pipeline("text-classification", model="j-hartmann/emotion-english-distilroberta-base", return_all_scores=True)
14
 
15
+ # Initialize NER pipeline
16
+ ner_tokenizer = AutoTokenizer.from_pretrained("dslim/bert-base-NER")
17
+ ner_model = AutoModelForTokenClassification.from_pretrained("dslim/bert-base-NER") # Renamed variable
18
+ ner_pipeline = pipeline("ner", model=ner_model, tokenizer=ner_tokenizer) # Renamed variable
19
 
20
  def allowed_file(filename):
21
  return '.' in filename and filename.rsplit('.', 1)[1].lower() in {'wav', 'mp3', 'ogg', 'flac', 'm4a'}
22
 
23
  @app.route('/transcribe', methods=['POST'])
24
  def transcribe_audio():
 
25
  if 'file' not in request.files:
26
  return jsonify({'error': 'No file uploaded'}), 400
27
 
28
  file = request.files['file']
29
 
 
30
  if file.filename == '':
31
  return jsonify({'error': 'No selected file'}), 400
32
 
 
33
  if not allowed_file(file.filename):
34
  return jsonify({'error': 'Unsupported file type'}), 400
35
 
36
  try:
 
37
  temp_path = "temp_audio"
38
  file.save(temp_path)
39
 
40
+ # Use the renamed whisper_model
41
+ result = whisper_model.transcribe(temp_path)
42
  transcription = result["text"]
43
 
 
44
  if os.path.exists(temp_path):
45
  os.remove(temp_path)
46
 
 
48
 
49
  except Exception as e:
50
  return jsonify({'error': str(e)}), 500
 
51
 
52
  @app.route('/classify', methods=['POST'])
53
  def classify():
 
63
  except Exception as e:
64
  return jsonify({"error": str(e)}), 500
65
 
 
 
 
66
  @app.route('/ner', methods=['POST'])
67
  def ner_endpoint():
68
  try:
 
69
  data = request.get_json()
70
  text = data.get("text", "")
71
 
72
+ # Use the renamed ner_pipeline
73
+ ner_results = ner_pipeline(text)
74
 
 
75
  words_and_entities = [
76
  {"word": result['word'], "entity": result['entity']}
77
  for result in ner_results
78
  ]
79
 
 
80
  return jsonify({"entities": words_and_entities})
81
  except Exception as e:
82
  return jsonify({"error": str(e)}), 500
83
+
84
+ if __name__ == '__main__':
85
+ app.run(host='0.0.0.0', port=5000, debug=True)