Update app.py

app.py

@@ -3,10 +3,12 @@
 Refactored Salama Assistant: text-only chatbot (STT and TTS removed)
 Drop this file into your Hugging Face Space (replace existing app.py) or run locally.
 
-This version:
-- Never passes device_map=None (avoids TypeError in accelerate)
-- Detects bitsandbytes availability and only requests 4-bit loading when safe
-- Keeps streaming responses into Gradio chat UI
+Performance-focused tweaks:
+- lower max_new_tokens
+- use greedy decoding (do_sample=False) for speed
+- call generate() under torch.no_grad()
+- set model.config.use_cache = True
+- other minor safe optimizations
 """
 
 import os
@@ -63,6 +65,10 @@ class WeeboAssistant:
             "You are an intelligent assistant. Answer questions briefly and accurately. "
             "Respond only in English. No long answers.\n"
         )
+        # set sensible defaults for generation speed
+        self.MAX_NEW_TOKENS = 256   # lowered from 512 for speed
+        self.DO_SAMPLE = False      # greedy = faster; set True if you need randomness
+        self.NUM_BEAMS = 1          # keep 1 for greedy; increase for beam search (slower)
         self._init_models()
 
     def _init_models(self):
@@ -82,6 +88,15 @@ class WeeboAssistant:
             self.llm_tokenizer = AutoTokenizer.from_pretrained(ADAPTER_REPO_ID, use_fast=True)
             print("Loaded tokenizer from ADAPTER_REPO_ID")
 
+        # ensure tokenizer has pad_token_id (some HF models lack it)
+        if getattr(self.llm_tokenizer, "pad_token_id", None) is None:
+            # try to set eos_token_id as pad if pad missing
+            if getattr(self.llm_tokenizer, "eos_token_id", None) is not None:
+                self.llm_tokenizer.pad_token_id = self.llm_tokenizer.eos_token_id
+            else:
+                # fallback to 0 (not ideal but prevents crashes)
+                self.llm_tokenizer.pad_token_id = 0
+
         if torch.cuda.is_available():
             device_map = "auto"
         else:
@@ -106,6 +121,11 @@ class WeeboAssistant:
                 BASE_MODEL_ID,
                 **base_model_kwargs,
             )
+            # make sure use_cache is enabled for faster autoregressive generation
+            try:
+                self.llm_model.config.use_cache = True
+            except Exception:
+                pass
             print("Base model loaded from", BASE_MODEL_ID)
         except Exception as e:
             raise RuntimeError(
@@ -132,6 +152,11 @@ class WeeboAssistant:
                 ADAPTER_REPO_ID,
                 **peft_kwargs,
             )
+            # ensure adapter-wrapped model also has use_cache
+            try:
+                self.llm_model.config.use_cache = True
+            except Exception:
+                pass
             print("PEFT adapter applied from", ADAPTER_REPO_ID)
         except Exception as e:
             raise RuntimeError(
@@ -165,27 +190,48 @@ class WeeboAssistant:
         prompt_lines.append("Assistant: ")
         prompt = "\n".join(prompt_lines)
 
-        inputs = self.llm_tokenizer(prompt, return_tensors="pt")
+        # Tokenize
+        inputs = self.llm_tokenizer(prompt, return_tensors="pt", padding=False)
         try:
             model_device = next(self.llm_model.parameters()).device
         except StopIteration:
             model_device = torch.device("cpu")
         inputs = {k: v.to(model_device) for k, v in inputs.items()}
 
+        # Streamer unchanged (still yields chunks)
         streamer = TextIteratorStreamer(self.llm_tokenizer, skip_prompt=True, skip_special_tokens=True)
 
+        # Prefill some generation kwargs optimized for speed
+        input_len = inputs["input_ids"].shape[1]
+        max_new = self.MAX_NEW_TOKENS
+        max_length = input_len + max_new
+
         generation_kwargs = dict(
             input_ids=inputs["input_ids"],
             attention_mask=inputs.get("attention_mask", None),
-            max_new_tokens=512,
-            do_sample=True,
-            temperature=0.6,
-            top_p=0.9,
+            max_length=max_length,               # prefer max_length = input_len + max_new_tokens
+            max_new_tokens=max_new,              # kept for clarity / compatibility
+            do_sample=self.DO_SAMPLE,            # greedy if False -> faster
+            num_beams=self.NUM_BEAMS,            # beam search >1 slows down; keep 1 for speed
             streamer=streamer,
             eos_token_id=getattr(self.llm_tokenizer, "eos_token_id", None),
+            pad_token_id=getattr(self.llm_tokenizer, "pad_token_id", None),
+            use_cache=True,
+            early_stopping=True,
         )
 
-        gen_thread = threading.Thread(target=self.llm_model.generate, kwargs=generation_kwargs, daemon=True)
+        # Run generate under no_grad for speed / memory
+        def _generate_thread():
+            with torch.no_grad():
+                try:
+                    # call generate on model (PEFT-wrapped)
+                    self.llm_model.generate(**generation_kwargs)
+                except Exception as e:
+                    # if streaming fails, put an error chunk into streamer by raising
+                    # streamer does not provide a direct API to inject text; print to log
+                    print("Generation error:", e)
+
+        gen_thread = threading.Thread(target=_generate_thread, daemon=True)
         gen_thread.start()
 
         return streamer