Update app.py

app.py

@@ -7,6 +7,10 @@ os.environ["HF_HUB_DISABLE_TELEMETRY"] = "1"
 os.environ["WANDB_DISABLED"] = "true"
 os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
 
+torch.set_num_threads(2)   # or 1 on the smallest CPU tier
+os.environ["OMP_NUM_THREADS"] = "2"
+os.environ["MKL_NUM_THREADS"] = "2"
+
 print("CUDA available:", torch.cuda.is_available())
 if torch.cuda.is_available():
     print("GPU:", torch.cuda.get_device_name(0), "| CC:", torch.cuda.get_device_capability(0))
@@ -89,7 +93,8 @@ def _nllb_translate(text: str, src_code: str, tgt_code: str, max_new: int = 256)
             forced_bos_token_id=tgt_id,
             max_new_tokens=max_new,
             do_sample=False,
-            num_beams=4,
+            # num_beams=4,
+            num_beams=1,
             repetition_penalty=1.02,
         )
     return mt_tok.batch_decode(out, skip_special_tokens=True)[0].strip()
@@ -170,7 +175,7 @@ def _restore_terms(s: str, placeholders: dict):
     return s
 
 @torch.inference_mode()
-def agri_answer_bullets_en(q_en: str, so_original: str, crop_hint: str = "", max_new=200) -> list[str]:
+def agri_answer_bullets_en(q_en: str, so_original: str, crop_hint: str = "", max_new=80) -> list[str]:
     crop = f" Crop: {crop_hint}." if crop_hint else ""
     prompt = f"<context> Somali original: {so_original}.{crop} {AGRI_HINT} <user> {q_en} <assistant>"
     enc = ag_tok(prompt, return_tensors="pt")
@@ -229,7 +234,8 @@ def answer_router(question_so: str) -> str:
     en_q_raw = so_to_en(question_so)
     en_q = polish_en_question(en_q_raw, crop)
     # Ask the agri model
-    bullets_en = agri_answer_bullets_en(en_q, question_so, crop_hint=crop, max_new=200)
+    # bullets_en = agri_answer_bullets_en(en_q, question_so, crop_hint=crop, max_new=200)
+    bullets_en = agri_answer_bullets_en(en_q, question_so, crop_hint=crop, max_new=80)
     if not bullets_en:
         return "Ma hubo. Fadlan ku celi su’aasha si kooban (dalagga, dhibaatada, iyo meesha)."
     # Back to Somali