Fix Phi-3 generate fallback on ZeroGPU

app.py

@@ -211,16 +211,38 @@ class GeneratorWrapper:
     def generate_stream(self, prompt: str, max_new_tokens: int, temperature: float, top_p: float):
         pipe = self.ensure()
         streamer = TextIteratorStreamer(pipe.tokenizer, skip_special_tokens=True, skip_prompt=True)
-        kwargs = {
+        inputs = pipe.tokenizer(prompt, return_tensors="pt")
+        device = getattr(pipe.model, "device", torch.device("cpu"))
+        inputs = {k: v.to(device) for k, v in inputs.items()}
+        gen_kwargs = {
             "max_new_tokens": max_new_tokens,
             "do_sample": True,
             "temperature": temperature,
             "top_p": top_p,
             "streamer": streamer,
-            "return_full_text": False,
+            "return_dict_in_generate": True,
+            "output_scores": False,
+            "use_cache": False,  # avoid DynamicCache issues on Phi-3 CPU
         }
-        thread = Thread(target=pipe, args=(prompt,), kwargs=kwargs)
-        thread.start()
+
+        def _run():
+            try:
+                pipe.model.generate(**inputs, **gen_kwargs)
+            except Exception as exc:
+                if self._fallback:
+                    print(f"[Generator:{self.name}] generate failed: {exc}; falling back to {self._fallback.name}")
+                    self._pipe = self._fallback.ensure()
+                    note = self._fallback_msg or f"Falling back to {self._fallback.name}."
+                    if note:
+                        streamer.put(note + " ")
+                    fb_stream = self._fallback.generate_stream(prompt, max_new_tokens, temperature, top_p)
+                    for tok in fb_stream:
+                        streamer.put(tok)
+                else:
+                    print(f"[Generator:{self.name}] generate failed: {exc}")
+                streamer.end()
+
+        Thread(target=_run, daemon=True).start()
         if self._note:
             yield self._note + " "
             self._note = None
@@ -243,13 +265,25 @@ def load_tinyllama():
 
 
 def load_phi3_mini():
-    return pipeline(
+    pipe = pipeline(
         "text-generation",
         model="microsoft/Phi-3-mini-4k-instruct",
         device_map="cpu",
         torch_dtype=torch.float32,
         trust_remote_code=True,
+        model_kwargs={
+            "use_cache": False,
+            "attn_implementation": "eager",
+        },
     )
+    # Disable cache to avoid DynamicCache.seen_tokens errors on ZeroGPU/CPU.
+    try:
+        pipe.model.config.use_cache = False
+        pipe.model.generation_config.use_cache = False
+        pipe.model.generation_config.cache_implementation = "static"
+    except Exception:
+        pass
+    return pipe
 
 
 _tiny_wrapper = GeneratorWrapper("tinyllama-1.1b-chat", load_tinyllama)
@@ -438,7 +472,7 @@ with gr.Blocks(title="MedDiscover") as demo:
             model_dd = gr.Dropdown(
                 label="Generator Model",
                 choices=list(GENERATORS.keys()),
-                value="phi-3-mini-4k",
+                value="tinyllama-1.1b-chat",
                 interactive=True,
             )
             k_slider = gr.Slider(1, 10, value=3, step=1, label="Top-k chunks")