AnyaSchen commited on
Commit
f6d1a77
·
1 Parent(s): 7cb050b

feat: use turbo model and return all probs

Browse files
Files changed (2) hide show
  1. language_detector.py +4 -2
  2. main.py +4 -3
language_detector.py CHANGED
@@ -27,6 +27,7 @@ class LanguageDetector:
27
  Returns:
28
  str: Detected language code (e.g., "en", "fr", etc.)
29
  float: Confidence score
 
30
  """
31
  try:
32
  # Load and preprocess audio
@@ -41,7 +42,7 @@ class LanguageDetector:
41
  detected_lang = max(probs, key=probs.get)
42
  confidence = probs[detected_lang]
43
 
44
- return detected_lang, confidence
45
 
46
  except Exception as e:
47
  logger.error(f"Error in language detection: {e}")
@@ -57,6 +58,7 @@ class LanguageDetector:
57
  Returns:
58
  str: Detected language code (e.g., "en", "fr", etc.)
59
  float: Confidence score
 
60
  """
61
  try:
62
  # Convert bytes to numpy array using librosa
@@ -77,7 +79,7 @@ class LanguageDetector:
77
  detected_lang = max(probs, key=probs.get)
78
  confidence = probs[detected_lang]
79
 
80
- return detected_lang, confidence
81
 
82
  except Exception as e:
83
  logger.error(f"Error in language detection: {e}")
 
27
  Returns:
28
  str: Detected language code (e.g., "en", "fr", etc.)
29
  float: Confidence score
30
+ dict: All language probabilities
31
  """
32
  try:
33
  # Load and preprocess audio
 
42
  detected_lang = max(probs, key=probs.get)
43
  confidence = probs[detected_lang]
44
 
45
+ return detected_lang, confidence, probs
46
 
47
  except Exception as e:
48
  logger.error(f"Error in language detection: {e}")
 
58
  Returns:
59
  str: Detected language code (e.g., "en", "fr", etc.)
60
  float: Confidence score
61
+ dict: All language probabilities
62
  """
63
  try:
64
  # Convert bytes to numpy array using librosa
 
79
  detected_lang = max(probs, key=probs.get)
80
  confidence = probs[detected_lang]
81
 
82
+ return detected_lang, confidence, probs
83
 
84
  except Exception as e:
85
  logger.error(f"Error in language detection: {e}")
main.py CHANGED
@@ -31,7 +31,7 @@ language_detector = None
31
  async def lifespan(app: FastAPI):
32
  global kit, language_detector
33
  kit = WhisperLiveKit()
34
- language_detector = LanguageDetector(model_name="medium")
35
  yield
36
 
37
  app = FastAPI(lifespan=lifespan)
@@ -68,14 +68,15 @@ async def detect_language(file: UploadFile = File(...)):
68
 
69
  # Use the language detector with the saved file
70
  if language_detector:
71
- detected_lang, confidence = language_detector.detect_language_from_file(file_path)
72
 
73
  # Clean up - remove the temporary file
74
  os.remove(file_path)
75
 
76
  return JSONResponse({
77
  "language": detected_lang,
78
- "confidence": float(confidence)
 
79
  })
80
  else:
81
  return JSONResponse(
 
31
  async def lifespan(app: FastAPI):
32
  global kit, language_detector
33
  kit = WhisperLiveKit()
34
+ language_detector = LanguageDetector(model_name="turbo")
35
  yield
36
 
37
  app = FastAPI(lifespan=lifespan)
 
68
 
69
  # Use the language detector with the saved file
70
  if language_detector:
71
+ detected_lang, confidence, probs = language_detector.detect_language_from_file(file_path)
72
 
73
  # Clean up - remove the temporary file
74
  os.remove(file_path)
75
 
76
  return JSONResponse({
77
  "language": detected_lang,
78
+ "confidence": float(confidence),
79
+ "probabilities": {lang: float(prob) for lang, prob in probs.items()}
80
  })
81
  else:
82
  return JSONResponse(