Retest models

src/recommend/main.py

@@ -3,5 +3,39 @@ from src.recommend.routes.ai_routes import router as ai_router
 from src.recommend.routes.system import router as system_router
 app = FastAPI()
 
-app.include_router(ai_router)
-app.include_router(system_router)
+from src.recommend.models.ai_request import Request
+import torch
+from transformers import BertForSequenceClassification, BertTokenizer
+
+model_name = "src/recommend/kalium_recommend"
+model = BertForSequenceClassification.from_pretrained(model_name)
+tokenizer = BertTokenizer.from_pretrained(model_name)
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model.to(device)
+
+
+@app.post("/ai")
+async def ai(request: Request):
+    text = (
+        f"This ads is focused in: {request.audience}"
+        f"This category is: {request.category} "
+        f"This area is: {request.area}"
+        f"This sub area of ads: {request.sub_area}"
+    )
+
+    inputs = tokenizer.encode_plus(
+        text,
+        add_special_tokens=True,
+        return_tensors="pt",
+        padding='max_length',
+        truncation=True,
+        max_length=255
+    )
+    input_ids = inputs['input_ids'].to(device)
+    attention_mask = inputs['attention_mask'].to(device)
+
+    with torch.no_grad():
+        outputs = model(input_ids, attention_mask=attention_mask)
+        logits = outputs.logits
+
+    return {"similarity": logits.tolist()}