Use GPU only for generate

app_zero.py

@@ -118,6 +118,48 @@ def finish_generation(text: str) -> str:
 
 
 @spaces.GPU
+def run_model(mode, title, text):
+    if mode == 0:
+        prompt = summarize_prompt(title, text)
+    elif mode == 50:
+        prompt = summarize_clickbait_short_prompt(title, text)
+    elif mode == 100:
+        prompt = summarize_clickbait_large_prompt(title, text)
+    else:
+        raise ValueError("Mode not supported")
+
+    formatted_prompt = tokenizer.apply_chat_template(
+        [{"role": "user", "content": prompt}],
+        tokenize=False,
+        add_generation_prompt=True,
+    )
+
+    model_inputs = tokenizer(
+        [formatted_prompt], return_tensors="pt", add_special_tokens=False
+    )
+
+    streamer = TextStreamer(
+        tokenizer=tokenizer, skip_prompt=True, skip_special_tokens=True
+    )
+
+    model_output = model.generate(
+        **model_inputs.to(model.device),
+        streamer=streamer,
+        generation_config=generation_config,
+        logits_processor=stop_criteria,
+    )
+
+    yield title, streamer, text
+
+    temp = tokenizer.batch_decode(
+        model_output[:, model_inputs["input_ids"].shape[-1] :],
+        skip_special_tokens=True,
+        clean_up_tokenization_spaces=True,
+    )[0]
+
+    return temp
+
+
 def generate_text(
     url: str, mode: int, progress=gr.Progress(track_tqdm=False)
 ) -> (str, str):
@@ -181,44 +223,7 @@ def generate_text(
         progress(0.5, desc="🤖 Leyendo noticia")
 
         try:
-            if mode == 0:
-                prompt = summarize_prompt(title, text)
-            elif mode == 50:
-                prompt = summarize_clickbait_short_prompt(title, text)
-            elif mode == 100:
-                prompt = summarize_clickbait_large_prompt(title, text)
-            else:
-                raise ValueError("Mode not supported")
-
-            formatted_prompt = tokenizer.apply_chat_template(
-                [{"role": "user", "content": prompt}],
-                tokenize=False,
-                add_generation_prompt=True,
-            )
-
-            model_inputs = tokenizer(
-                [formatted_prompt], return_tensors="pt", add_special_tokens=False
-            )
-
-            streamer = TextStreamer(
-                tokenizer=tokenizer, skip_prompt=True, skip_special_tokens=True
-            )
-
-            model_output = model.generate(
-                **model_inputs.to(model.device),
-                streamer=streamer,
-                generation_config=generation_config,
-                logits_processor=stop_criteria,
-            )
-
-            yield title, streamer, text
-
-            temp = tokenizer.batch_decode(
-                model_output[:, model_inputs["input_ids"].shape[-1] :],
-                skip_special_tokens=True,
-                clean_up_tokenization_spaces=True,
-            )[0]
-
+            temp = run_model(mode, title, text)
             yield title, temp, text
 
         except Exception as e: