-
Notifications
You must be signed in to change notification settings - Fork 0
/
main.py
68 lines (54 loc) · 2.27 KB
/
main.py
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
61
62
63
64
65
66
67
68
import fasttext
import numpy as np
from fastapi import FastAPI
from classes import LanguageDetectionResponse
from classes import RequestBody
from classes import ResponseBody
# Specify the number of other languages to detect
NUM_OTHER_LANGUAGES = 3
fasttext.FastText.eprint = lambda x: None
# Load the FastText model
model_path = "models/model.bin"
model = fasttext.load_model(model_path)
app = FastAPI()
# Handler function for language detection
@app.post("/detect")
async def detect_language(request: RequestBody) -> LanguageDetectionResponse:
# Perform language detection using the pre-loaded FastText model
predictions = model.predict(request.text, k=NUM_OTHER_LANGUAGES)
if predictions:
# Extract the primary language prediction
primary_predictions = predictions[0]
primary_probabilities = predictions[1]
# Clean language labels
primary_predictions = [
pred.replace("__label__", "") for pred in primary_predictions
]
# Sort the primary predictions and probabilities
sorted_indices = np.argsort(primary_probabilities)[::-1]
primary_predictions = np.array(primary_predictions)[sorted_indices]
primary_probabilities = primary_probabilities[sorted_indices]
# Extract the top prediction
primary_prediction = primary_predictions[0]
primary_prediction_probability = primary_probabilities[0]
# Extract the other predictions
other_languages = [
ResponseBody(language=pred, accuracy=prob)
for pred, prob in zip(
primary_predictions[: NUM_OTHER_LANGUAGES + 1],
primary_probabilities[: NUM_OTHER_LANGUAGES + 1],
)
]
# Construct the response with language detection results
response = LanguageDetectionResponse(
request_text=request.text,
primary_language=ResponseBody(
language=primary_prediction, accuracy=primary_prediction_probability
),
other_languages=other_languages,
)
# Return a successful response with the language detection results
return response
else:
# Return an internal server error if language detection fails
return {"error": "Failed to detect language"}