Update app.py

app.py

@@ -2,19 +2,18 @@
 
 import streamlit as st
 import logging
+import torch
 import pandas as pd
 import plotly.express as px
-import torch
-from typing import Union, List
-
-from langdetect import detect, LangDetectException
 from transformers import (
     AutoTokenizer,
     AutoModelForSeq2SeqLM,
     pipeline,
     BitsAndBytesConfig,
 )
+from langdetect import detect, LangDetectException
 import evaluate
+from typing import Union, List
 
 # ────────── Logging ──────────
 logging.basicConfig(
@@ -38,7 +37,7 @@ class ModelManager:
         quantize: bool = True,
         default_tgt: str = None,
     ):
-        # If user requested quantization but CUDA isn't available, disable it
+        # Disable 8-bit if CUDA isn't available
         if quantize and not torch.cuda.is_available():
             logger.warning("CUDA unavailable; disabling 8-bit quantization")
             quantize = False
@@ -50,7 +49,7 @@ class ModelManager:
         ]
         self.default_tgt = default_tgt  # will auto-pick if None
 
-        self.selected_model_name: str = None
+        self.model_name: str = None
         self.tokenizer = None
         self.model = None
         self.pipeline = None
@@ -66,14 +65,10 @@ class ModelManager:
                 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"):
-                    raise AttributeError(
-                        f"Tokenizer for {model_name} missing lang_code_to_id"
-                    )
+                    raise AttributeError(f"Tokenizer for {model_name} missing lang_code_to_id")
 
                 # Load model (with or without 8-bit)
-                logger.info(
-                    f"Loading model {model_name} (8-bit={self.quantize})"
-                )
+                logger.info(f"Loading model {model_name} (8-bit={self.quantize})")
                 if self.quantize:
                     bnb_cfg = BitsAndBytesConfig(load_in_8bit=True)
                     mdl = AutoModelForSeq2SeqLM.from_pretrained(
@@ -96,7 +91,7 @@ class ModelManager:
                 )
 
                 # Store and break
-                self.selected_model_name = model_name
+                self.model_name = model_name
                 self.tokenizer = tok
                 self.model = mdl
                 self.pipeline = pipe
@@ -104,22 +99,19 @@ class ModelManager:
 
                 # Auto-pick Turkish target code if none specified
                 if not self.default_tgt:
-                    tur_codes = [
-                        c for c in self.lang_codes if c.lower().startswith("tr")
-                    ]
+                    tur_codes = [c for c in self.lang_codes if c.lower().startswith("tr")]
                     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 {self.candidates}: {last_err}"
-        )
+        raise RuntimeError(f"Could not load any model from {self.candidates}: {last_err}")
 
     def translate(
         self,
@@ -138,9 +130,7 @@ class ModelManager:
             sample = text[0] if isinstance(text, list) else text
             try:
                 iso = detect(sample).lower()
-                candidates = [
-                    c for c in self.lang_codes if c.lower().startswith(iso)
-                ]
+                candidates = [c for c in self.lang_codes if c.lower().startswith(iso)]
                 if not candidates:
                     raise LangDetectException(f"No code for ISO '{iso}'")
                 exact = [c for c in candidates if c.lower() == iso]
@@ -148,9 +138,7 @@ class ModelManager:
                 logger.info(f"Auto-detected src_lang={src}")
             except Exception as e:
                 logger.warning(f"langdetect failed ({e}); defaulting to English")
-                eng_codes = [
-                    c for c in self.lang_codes if c.lower().startswith("en")
-                ]
+                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
@@ -159,13 +147,14 @@ class ModelManager:
 
     def get_info(self):
         """Return metadata for the sidebar display."""
+        quantized = getattr(self.model, "is_loaded_in_8bit", False)
         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"
+            dev = self.model.device
+            device = str(dev) if isinstance(dev, torch.device) else f"cuda:{getattr(dev, 'index', '')}"
         return {
-            "model":       self.selected_model_name,
-            "quantized":   self.quantize,
+            "model":       self.model_name,
+            "quantized":   quantized,
             "device":      device,
             "default_tgt": self.default_tgt,
         }
@@ -175,9 +164,16 @@ class ModelManager:
 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")
-        logger.info("Loaded BLEU, BERTScore, COMET metrics")
+        try:
+            self.bertscore = evaluate.load("bertscore")
+        except Exception as e:
+            logger.error("BERTScore load error: %s", e)
+            self.bertscore = None
+        try:
+            self.comet = evaluate.load("comet", model_id="unbabel/comet-mqm-qe-da")
+        except Exception as e:
+            logger.error("COMET load error: %s", e)
+            self.comet = None
 
     def evaluate(
         self,
@@ -186,36 +182,71 @@ class TranslationEvaluator:
         predictions: List[str],
     ):
         results = {}
+
         # BLEU
-        results["BLEU"] = self.bleu.compute(
-            predictions=predictions,
-            references=[[r] for r in references],
-        )["bleu"]
+        try:
+            bleu_res = self.bleu.compute(
+                predictions=predictions,
+                references=[[r] for r in references],
+            )
+            results["BLEU"] = float(bleu_res.get("bleu", 0.0))
+        except Exception as e:
+            logger.error("BLEU compute error: %s", e)
+            results["BLEU"] = 0.0
+
         # 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
+        if self.bertscore:
+            try:
+                bs = self.bertscore.compute(
+                    predictions=predictions, references=references, lang="xx"
+                )
+                f1 = bs.get("f1", [])
+                results["BERTScore"] = float(sum(f1)) / max(len(f1), 1)
+            except Exception as e:
+                logger.error("BERTScore compute error: %s", e)
+                results["BERTScore"] = 0.0
+        else:
+            results["BERTScore"] = 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
+        if self.bertscore:
+            try:
+                bs_tr = self.bertscore.compute(
+                    predictions=predictions, references=references, lang="tr"
+                )
+                f1_tr = bs_tr.get("f1", [])
+                results["BERTurk"] = float(sum(f1_tr)) / max(len(f1_tr), 1)
+            except Exception as e:
+                logger.error("BERTurk compute error: %s", e)
+                results["BERTurk"] = 0.0
+        else:
+            results["BERTurk"] = 0.0
+
         # COMET
-        cm = self.comet.compute(
-            srcs=sources, hyps=predictions, refs=references
-        )
-        scores = cm.get("scores", None)
-        results["COMET"] = float(scores[0] if isinstance(scores, list) else scores) or 0.0
+        if self.comet:
+            try:
+                cm = self.comet.compute(
+                    srcs=sources, hyps=predictions, refs=references
+                )
+                sc = cm.get("scores", None)
+                if isinstance(sc, list):
+                    results["COMET"] = float(sc[0]) if sc else 0.0
+                else:
+                    results["COMET"] = float(sc or 0.0)
+            except Exception as e:
+                logger.error("COMET compute error: %s", e)
+                results["COMET"] = 0.0
+        else:
+            results["COMET"] = 0.0
+
         return results
 
 
 # ────────── Streamlit App ──────────
-
 @st.cache_resource
 def load_resources():
     mgr = ModelManager(quantize=True)
-    ev  = TranslationEvaluator()
+    ev = TranslationEvaluator()
     return mgr, ev
 
 
@@ -235,17 +266,17 @@ def process_text(
     metrics: List[str],
 ):
     out = mgr.translate(src)
-    hyp = out[0]["translation_text"]
+    hyp = out[0]["translation_text"] if isinstance(out, list) else out["translation_text"]
     scores = ev.evaluate([src], [ref or ""], [hyp])
     return {
         "source":     src,
         "reference":  ref,
         "hypothesis": hyp,
-        **{m: scores[m] for m in metrics},
+        **{m: scores.get(m, 0.0) for m in metrics},
     }
 
 
-def _show_single_results(res: dict):
+def show_single_results(res: dict, metrics: List[str]):
     left, right = st.columns(2)
     with left:
         st.markdown("**Source:**")
@@ -257,7 +288,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 metrics}])
+        df = pd.DataFrame([{m: res[m] for m in metrics}])
         st.table(df)
 
 
@@ -283,13 +314,13 @@ def process_file(
         for s, r, h in zip(srcs, refs, hyps):
             sc = ev.evaluate([s], [r], [h])
             entry = {"src": s, "ref_tr": r, "hyp_tr": h}
-            entry.update({m: sc[m] for m in metrics})
+            entry.update({m: sc.get(m, 0.0) for m in metrics})
             results.append(entry)
         prog.progress(min(i + batch_size, total) / total)
     return pd.DataFrame(results)
 
 
-def _show_batch_viz(df: pd.DataFrame, metrics: List[str]):
+def show_batch_viz(df: pd.DataFrame, metrics: List[str]):
     for m in metrics:
         st.markdown(f"#### {m} Distribution")
         fig = px.histogram(df, x=m)
@@ -297,9 +328,7 @@ def _show_batch_viz(df: pd.DataFrame, metrics: List[str]):
 
 
 def main():
-    st.set_page_config(
-        page_title="🔤 Translation→Turkish Quality", layout="wide"
-    )
+    st.set_page_config(page_title="🔤 Translation→Turkish Quality", layout="wide")
     st.title("🔤 Translation → TR Quality & COMET")
     st.markdown(
         "Translate any language into Turkish and evaluate with BLEU, BERTScore, BERTurk & COMET."
@@ -326,7 +355,7 @@ def main():
         if st.button("Evaluate"):
             with st.spinner("Translating & evaluating…"):
                 res = process_text(src, ref, mgr, ev, metrics)
-            _show_single_results(res)
+            show_single_results(res, metrics)
 
     with tab2:
         uploaded = st.file_uploader(
@@ -337,10 +366,8 @@ def main():
                 df_res = process_file(uploaded, mgr, ev, metrics, batch_size)
             st.markdown("### Batch Results")
             st.dataframe(df_res, use_container_width=True)
-            _show_batch_viz(df_res, metrics)
-            st.download_button(
-                "Download CSV", df_res.to_csv(index=False), "results.csv"
-            )
+            show_batch_viz(df_res, metrics)
+            st.download_button("Download CSV", df_res.to_csv(index=False), "results.csv")
 
 
 if __name__ == "__main__":