Update app.py

app.py

@@ -2,12 +2,14 @@ from fastapi import FastAPI, Request
 from transformers import AutoModelForSequenceClassification, AutoTokenizer, pipeline
 import torch
 from pydantic import BaseModel
+from typing import Optional
 
 app = FastAPI()
 
 
 class InputText(BaseModel):
     text : str
+    threshold: Optional[float] = None
 
 
 model_name = "cardiffnlp/twitter-xlm-roberta-base-sentiment"
@@ -15,11 +17,45 @@ sentiment_model = AutoModelForSequenceClassification.from_pretrained(model_name)
 sentiment_tokenizer = AutoTokenizer.from_pretrained(model_name)
 sentiment_model.config.id2label[3] = "mixed"
 
+model_name = 'qanastek/51-languages-classifier'
+language_model = AutoModelForSequenceClassification.from_pretrained(model_name)
+language_tokenizer = AutoTokenizer.from_pretrained(model_name)
+
+
+
+
+
+
+
 @app.get("/")
 def greet_json():
     return {"Hello": "World!"}
 
 
+
+@app.post("/language_detection")
+async def sentiment_score(inp: InputText):
+    inputs = tokenizer(inp.text, return_tensors='pt')
+    with torch.no_grad():
+        logits = language_model(**inputs).logits
+    
+    softmax = torch.nn.functional.sigmoid(logits)
+    
+    # Apply the threshold by creating a mask
+    mask = softmax >= inp.threshold
+    
+    # Filter the tensor based on the threshold
+    filtered_x = softmax[mask]
+    
+    # Get the sorted indices of the filtered tensor
+    sorted_indices = torch.argsort(filtered_x, descending=True)
+    
+    # Map the sorted indices back to the original tensor indices
+    original_indices = torch.nonzero(mask, as_tuple=True)[1][sorted_indices]
+    
+    return [{"label":model.config.id2label[predicted_class_id.tolist()], "score":softmax[0, predicted_class_id].tolist()} for predicted_class_id in original_indices]
+
+
 @app.post("/sentiment_score")
 async def sentiment_score(inp: InputText):
     text = inp.text