Update app.py

app.py

@@ -1,7 +1,10 @@
+# app.py
+
 import streamlit as st
 import logging
 import pandas as pd
 import plotly.express as px
+import torch
 from typing import Union, List
 
 from langdetect import detect, LangDetectException
@@ -22,36 +25,44 @@ logging.basicConfig(
 logger = logging.getLogger(__name__)
 
 
-# ────────── Model Management ──────────
+# ────────── Model Manager ──────────
 class ModelManager:
     """
-    Automatically selects, loads, and wraps a seq2seq translation model
-    in 8-bit (with FP32 fallback), plus language‐code auto-detection.
+    Selects and loads a translation model (NLLB-200 or M2M100),
+    using 8-bit quantization only if CUDA is available.
+    Auto-detects source language and defaults target to Turkish.
     """
-
     def __init__(
         self,
         candidates: List[str] = None,
         quantize: bool = True,
         default_tgt: str = None,
     ):
+        # If user requested quantization but CUDA isn't available, disable it
+        if quantize and not torch.cuda.is_available():
+            logger.warning("CUDA unavailable; disabling 8-bit quantization")
+            quantize = False
+        self.quantize = quantize
+
         self.candidates = candidates or [
             "facebook/nllb-200-distilled-600M",
             "facebook/m2m100_418M",
         ]
-        self.quantize = quantize
-        self.default_tgt = default_tgt  # if None → auto-pick Turkish
+        self.default_tgt = default_tgt  # will auto-pick if None
+
+        self.selected_model_name: str = None
         self.tokenizer = None
         self.model = None
         self.pipeline = None
         self.lang_codes: List[str] = []
+
         self._select_and_load()
 
     def _select_and_load(self):
         last_err = None
         for model_name in self.candidates:
             try:
-                # 1) Load tokenizer
+                # Load tokenizer
                 logger.info(f"Loading tokenizer for {model_name}")
                 tok = AutoTokenizer.from_pretrained(model_name, use_fast=True)
                 if not hasattr(tok, "lang_code_to_id"):
@@ -59,53 +70,55 @@ class ModelManager:
                         f"Tokenizer for {model_name} missing lang_code_to_id"
                     )
 
-                # 2) Load model with bitsandbytes 8-bit quantization
+                # Load model (with or without 8-bit)
                 logger.info(
-                    f"Loading model {model_name} "
-                    f"(8-bit={'on' if self.quantize else 'off'})"
-                )
-                bnb_cfg = BitsAndBytesConfig(load_in_8bit=self.quantize)
-                model = AutoModelForSeq2SeqLM.from_pretrained(
-                    model_name,
-                    device_map="auto",
-                    quantization_config=bnb_cfg,
+                    f"Loading model {model_name} (8-bit={self.quantize})"
                 )
+                if self.quantize:
+                    bnb_cfg = BitsAndBytesConfig(load_in_8bit=True)
+                    mdl = AutoModelForSeq2SeqLM.from_pretrained(
+                        model_name,
+                        device_map="auto",
+                        quantization_config=bnb_cfg,
+                    )
+                else:
+                    mdl = AutoModelForSeq2SeqLM.from_pretrained(
+                        model_name,
+                        device_map="auto",
+                    )
                 logger.info(f"Model {model_name} loaded successfully")
 
-                # 3) Build a translation pipeline around it
+                # Wrap in a translation pipeline
                 pipe = pipeline(
                     "translation",
-                    model=model,
+                    model=mdl,
                     tokenizer=tok,
                 )
 
-                # 4) On success, store and break
+                # Store and break
+                self.selected_model_name = model_name
                 self.tokenizer = tok
-                self.model = model
+                self.model = mdl
                 self.pipeline = pipe
                 self.lang_codes = list(tok.lang_code_to_id.keys())
-                logger.info(f"Available language codes: {self.lang_codes[:5]}…")
 
-                # 5) Auto-pick Turkish target if needed
+                # Auto-pick Turkish target code if none specified
                 if not self.default_tgt:
-                    tur = [
-                        code
-                        for code in self.lang_codes
-                        if code.lower().startswith("tr")
+                    tur_codes = [
+                        c for c in self.lang_codes if c.lower().startswith("tr")
                     ]
-                    if not tur:
-                        raise ValueError(f"No Turkish code in {model_name}")
-                    self.default_tgt = tur[0]
+                    if not tur_codes:
+                        raise ValueError(f"No Turkish code found in {model_name}")
+                    self.default_tgt = tur_codes[0]
                 logger.info(f"Default target language: {self.default_tgt}")
 
                 return
-
             except Exception as e:
                 logger.warning(f"Failed to load {model_name}: {e}")
                 last_err = e
 
         raise RuntimeError(
-            f"Could not load any model from candidates {self.candidates}: {last_err}"
+            f"Could not load any model from {self.candidates}: {last_err}"
         )
 
     def translate(
@@ -116,13 +129,11 @@ class ModelManager:
     ):
         """
         Translate `text` from src_lang → tgt_lang.
-        If src_lang is None: auto-detect via langdetect.
-        If tgt_lang is None: use default_tgt (Turkish).
-        Returns the pipeline output (list of dicts with 'translation_text').
+        Auto-detects src_lang if not given.
         """
         tgt = tgt_lang or self.default_tgt
 
-        # Auto-detect source
+        # Auto-detect source language if missing
         if not src_lang:
             sample = text[0] if isinstance(text, list) else text
             try:
@@ -132,41 +143,41 @@ class ModelManager:
                 ]
                 if not candidates:
                     raise LangDetectException(f"No code for ISO '{iso}'")
-                # prefer exact match
                 exact = [c for c in candidates if c.lower() == iso]
                 src = exact[0] if exact else candidates[0]
                 logger.info(f"Auto-detected src_lang={src}")
             except Exception as e:
                 logger.warning(f"langdetect failed ({e}); defaulting to English")
-                eng = [c for c in self.lang_codes if c.lower().startswith("en")]
-                src = eng[0] if eng else self.lang_codes[0]
+                eng_codes = [
+                    c for c in self.lang_codes if c.lower().startswith("en")
+                ]
+                src = eng_codes[0] if eng_codes else self.lang_codes[0]
         else:
             src = src_lang
 
-        # Call the pipeline with both src_lang and tgt_lang
         return self.pipeline(text, src_lang=src, tgt_lang=tgt)
 
     def get_info(self):
-        """Return metadata for sidebar display."""
-        model = getattr(self.model, "config", None)
-        quantized = getattr(self.model, "is_loaded_in_8bit", False)
-        device = getattr(self.model.device, "index", None)
-        device = f"cuda:{device}" if device is not None else "cpu"
+        """Return metadata for the sidebar display."""
+        device = "cpu"
+        if torch.cuda.is_available() and hasattr(self.model, "device"):
+            idx = self.model.device.index if hasattr(self.model.device, "index") else None
+            device = f"cuda:{idx}" if idx is not None else "cuda"
         return {
-            "model":     self.model.name_or_path,
-            "quantized": quantized,
-            "device":    device,
+            "model":       self.selected_model_name,
+            "quantized":   self.quantize,
+            "device":      device,
             "default_tgt": self.default_tgt,
         }
 
 
-# ────────── Evaluation ──────────
+# ────────── Evaluator ──────────
 class TranslationEvaluator:
     def __init__(self):
         self.bleu = evaluate.load("bleu")
         self.bertscore = evaluate.load("bertscore")
         self.comet = evaluate.load("comet", model_id="unbabel/comet-mqm-qe-da")
-        logging.info("Loaded BLEU, BERTScore, COMET")
+        logger.info("Loaded BLEU, BERTScore, COMET metrics")
 
     def evaluate(
         self,
@@ -175,36 +186,27 @@ class TranslationEvaluator:
         predictions: List[str],
     ):
         results = {}
-
         # BLEU
         results["BLEU"] = self.bleu.compute(
             predictions=predictions,
             references=[[r] for r in references],
         )["bleu"]
-
         # BERTScore (general)
         bs = self.bertscore.compute(
             predictions=predictions, references=references, lang="xx"
         )
         results["BERTScore"] = sum(bs["f1"]) / len(bs["f1"]) if bs["f1"] else 0.0
-
         # BERTurk (Turkish)
         bs_tr = self.bertscore.compute(
             predictions=predictions, references=references, lang="tr"
         )
         results["BERTurk"] = sum(bs_tr["f1"]) / len(bs_tr["f1"]) if bs_tr["f1"] else 0.0
-
         # COMET
-        co = self.comet.compute(
+        cm = self.comet.compute(
             srcs=sources, hyps=predictions, refs=references
         )
-        # `scores` may be a float or list
-        score = co.get("scores", None)
-        if isinstance(score, list):
-            results["COMET"] = score[0] if score else 0.0
-        else:
-            results["COMET"] = score or 0.0
-
+        scores = cm.get("scores", None)
+        results["COMET"] = float(scores[0] if isinstance(scores, list) else scores) or 0.0
         return results
 
 
@@ -212,9 +214,6 @@ class TranslationEvaluator:
 
 @st.cache_resource
 def load_resources():
-    """
-    Load and cache ModelManager & TranslationEvaluator on first run.
-    """
     mgr = ModelManager(quantize=True)
     ev  = TranslationEvaluator()
     return mgr, ev
@@ -235,7 +234,7 @@ def process_text(
     ev: TranslationEvaluator,
     metrics: List[str],
 ):
-    out = mgr.translate(src)  # list of dicts
+    out = mgr.translate(src)
     hyp = out[0]["translation_text"]
     scores = ev.evaluate([src], [ref or ""], [hyp])
     return {
@@ -258,7 +257,7 @@ def _show_single_results(res: dict):
             st.write(res["reference"])
     with right:
         st.markdown("### Scores")
-        df = pd.DataFrame([{k: v for k, v in res.items() if k in res.keys() and k in ["BLEU","BERTScore","BERTurk","COMET"]}])
+        df = pd.DataFrame([{k: v for k, v in res.items() if k in metrics}])
         st.table(df)
 
 
@@ -279,7 +278,7 @@ def process_file(
         batch = df.iloc[i : i + batch_size]
         srcs = batch["src"].tolist()
         refs = batch["ref_tr"].tolist()
-        outs = mgr.translate(srcs)  # batch translation
+        outs = mgr.translate(srcs)
         hyps = [o["translation_text"] for o in outs]
         for s, r, h in zip(srcs, refs, hyps):
             sc = ev.evaluate([s], [r], [h])