Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,3 +1,4 @@
+# streamlit_app.py
 import time
 import torch
 import streamlit as st
@@ -22,16 +23,14 @@ GENERATION_CONFIG: Dict[str, Any] = {
     "repetition_penalty": 1.10, "seed": 42, "sort_by": "logp/len",
 }
 
-# ------------------ MODEL LOADING (LOCAL 4-BIT) ------------------
+# ------------------ MODEL LOADING (CPU/GPU AUTO) ------------------
 @st.cache_resource
 def load_model() -> Tuple[MT5Tokenizer, MT5ForConditionalGeneration, torch.device]:
+    # Avoid CUDA initialization if no driver; select device explicitly.
+    device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
     tok = MT5Tokenizer.from_pretrained(MODEL_PATH, cache_dir=CACHE_DIR)
-    model = MT5ForConditionalGeneration.from_pretrained(
-        MODEL_PATH,
-        cache_dir=CACHE_DIR,
-        device_map={"": 0}
-    )
-    device = model.device
+    model = MT5ForConditionalGeneration.from_pretrained(MODEL_PATH, cache_dir=CACHE_DIR)
+    model.to(device)
     return tok, model, device
 
 # ------------------ GENERATION HELPERS ------------------
@@ -44,8 +43,7 @@ def decode_sequences(tok: MT5Tokenizer, seqs: torch.Tensor) -> List[str]:
     return tok.batch_decode(seqs, skip_special_tokens=True)
 
 def avg_logprobs_from_generate(tok: MT5Tokenizer, gen) -> List[float]:
-    if not hasattr(gen, "scores"):
-        return [float("nan")] * gen.sequences.size(0)
+    if not hasattr(gen, "scores"): return [float("nan")] * gen.sequences.size(0)
     scores, seqs = gen.scores, gen.sequences
     nseq, eos_id, pad_id = seqs.size(0), tok.eos_token_id or 1, tok.pad_token_id
     sum_logp = torch.zeros(nseq, dtype=torch.float32, device=scores[0].device)
@@ -66,27 +64,15 @@ def avg_logprobs_from_generate(tok: MT5Tokenizer, gen) -> List[float]:
 # --- UPDATED sampling_generate function ---
 def sampling_generate(tok, model, device, inputs, top_n, temperature, top_p, no_repeat_ngram_size, repetition_penalty, bad_words_ids: List[List[int]] = None):
     """Now accepts a list of 'bad_words_ids' to forbid certain sequences."""
-    kwargs = dict(
-        max_length=MAX_TARGET_LENGTH,
-        do_sample=True,
-        temperature=temperature,
-        top_p=top_p,
-        num_return_sequences=top_n,
-        return_dict_in_generate=True,
-        output_scores=True
-    )
-    if no_repeat_ngram_size > 0:
-        kwargs["no_repeat_ngram_size"] = int(no_repeat_ngram_size)
-    if repetition_penalty != 1.0:
-        kwargs["repetition_penalty"] = float(repetition_penalty)
-    if bad_words_ids:
-        kwargs["bad_words_ids"] = bad_words_ids
+    kwargs = dict(max_length=MAX_TARGET_LENGTH, do_sample=True, temperature=temperature, top_p=top_p, num_return_sequences=top_n, return_dict_in_generate=True, output_scores=True)
+    if no_repeat_ngram_size > 0: kwargs["no_repeat_ngram_size"] = int(no_repeat_ngram_size)
+    if repetition_penalty != 1.0: kwargs["repetition_penalty"] = float(repetition_penalty)
+    if bad_words_ids: kwargs["bad_words_ids"] = bad_words_ids
 
     gen = model.generate(**inputs, **kwargs)
     return decode_sequences(tok, gen.sequences), avg_logprobs_from_generate(tok, gen)
 
-def normalize_text(s: str) -> str:
-    return " ".join(s.strip().lower().split())
+def normalize_text(s: str) -> str: return " ".join(s.strip().lower().split())
 
 # ------------------ STREAMLIT APP ------------------
 st.set_page_config(
@@ -126,23 +112,16 @@ if run_button:
         seen_texts_for_bad_words = set()
 
         num_batches = (TOTAL_DESIRED_CANDIDATES + GENERATION_BATCH_SIZE - 1) // GENERATION_BATCH_SIZE
-        
         progress_bar = st.progress(0)
 
         for i in range(num_batches):
             current_seed = cfg["seed"] + i
             torch.manual_seed(current_seed)
-            if torch.cuda.is_available():
-                torch.cuda.manual_seed_all(current_seed)
+            if torch.cuda.is_available(): torch.cuda.manual_seed_all(current_seed)
 
             bad_words_ids = None
             if seen_texts_for_bad_words:
-                bad_words_ids = tok(
-                    list(seen_texts_for_bad_words),
-                    add_special_tokens=False,
-                    padding=True,
-                    truncation=True
-                )["input_ids"]
+                bad_words_ids = tok(list(seen_texts_for_bad_words), add_special_tokens=False, padding=True, truncation=True)["input_ids"]
 
             batch_texts, batch_scores = sampling_generate(
                 tok, model, device, inputs,
@@ -157,11 +136,11 @@ if run_button:
             all_texts.extend(batch_texts)
             all_scores.extend(batch_scores)
             for txt in batch_texts:
-                if txt:
-                    seen_texts_for_bad_words.add(txt)
+                if txt: seen_texts_for_bad_words.add(txt)
             
             progress_bar.progress((i + 1) / num_batches)
 
+        # Deduplicate and finalize the list
         final_enriched = []
         final_seen_normalized = set()
         for txt, sc in zip(all_texts, all_scores):