Fix bug

app.py

@@ -1,16 +1,72 @@
 import os
 
 import gradio as gr
+import copy
+from llama_cpp import Llama
+
+# CMAKE_ARGS="-DLLAMA_CUBLAS=on" FORCE_CMAKE=1 pip install llama-cpp-python --no-cache-dir
+# CMAKE_ARGS="-DLLAMA_BLAS=ON -DLLAMA_BLAS_VENDOR=OpenBLAS" FORCE_CMAKE=1 pip install llama-cpp-python --no-cache-dir
+
+import json
+import datetime
+from transformers import AutoTokenizer
 from download_url import download_text_and_title
+from prompts import clickbait_prompt, summary_prompt, clickbait_summary_prompt
 from cache_system import CacheHandler
-from gradio_client import Client
+from huggingface_hub import hf_hub_download
+
+auth_token = os.environ.get("TOKEN_FROM_SECRET") or True
 
 print(f"CPU cores: {os.cpu_count()}.")
 
-server = os.environ.get("SERVER") or True
-auth_token = os.environ.get("TOKEN") or True
+llm = Llama(
+    model_path=hf_hub_download(
+        repo_id=os.environ.get("REPO_ID", "Iker/ClickbaitFighter-10B"),
+        filename=os.environ.get("MODEL_FILE", "ClickbaitFighter-10B_q4_k_m.gguf"),
+        token=auth_token,
+    ),
+    n_ctx=0,
+    n_gpu_layers=-1,  # change n_gpu_layers if you have more or less VRAM
+    n_threads=8,
+)
+
+tokenizer = AutoTokenizer.from_pretrained(
+    "Iker/ClickbaitFighter-10B",
+    add_eos_token=True,
+    token=auth_token,
+    use_fast=True,
+)
 
-client = Client(server)
+
+def generate_prompt(
+    tittle: str,
+    body: str,
+    mode: str = "finetune",
+) -> str:
+    """
+    Generate the prompt for the model.
+
+    Args:
+        tittle (`str`):
+            The tittle of the article.
+        body (`str`):
+            The body of the article.
+        mode (`str`):
+            The mode of the model. Can be 'clickbait', 'summary' or 'clickbait-summary'.
+    Returns:
+        `str`: The formatted prompt.
+    """
+
+    if mode == "clickbait":
+        return clickbait_prompt(tittle, body)
+    elif mode == "summary":
+        return summary_prompt(tittle, body)
+    elif mode == "clickbait-summary":
+        return clickbait_summary_prompt(tittle, body)
+    else:
+        raise ValueError(
+            "Invalid mode. Valid modes are 'clickbait', 'summary' and 'clickbait-summary'"
+        )
 
 
 def finish_generation(text: str) -> str:
@@ -41,40 +97,52 @@ def generate_text(
         if title is None or text is None:
             yield (
                 "🤖 No he podido acceder a la notica, asegurate que la URL es correcta y que es posible acceder a la noticia desde un navegador.",
-                "❌❌❌ Inténtalo de nuevo ❌❌❌",
+                "❌❌❌",
                 "Error",
             )
             return (
                 "🤖 No he podido acceder a la notica, asegurate que la URL es correcta y que es posible acceder a la noticia desde un navegador.",
-                "❌❌❌ Inténtalo de nuevo ❌❌❌",
+                "❌❌❌",
                 "Error",
             )
 
         progress(0.5, desc="🤖 Leyendo noticia")
 
-        try:
-            temp = client.predict(
-                url,  # str  in '🌐 URL de la noticia' Textbox component
-                title,  # str  in '🌐 Título de la noticia' Textbox component
-                text,  # str  in '📰 Cuerpo de la noticia' Textbox component
-                mode,  # float (numeric value between 0 and 100) in '🎚️ Nivel de resumen' Slider component
-                api_name="/predict",
-            )
+        # 2) Generate the prompt
+        if mode == 0:
+            mo = "summary"
+        elif mode == 100:
+            mo = "clickbait"
+        else:
+            mo = "clickbait-summary"
+        input_prompt = generate_prompt(title, text, mo)
+        input_prompt = tokenizer.apply_chat_template(
+            [{"role": "user", "content": input_prompt}],
+            tokenize=False,
+            add_generation_prompt=True,
+        )
 
-            for o in temp:
-                yield title, o, text
+        output = llm(
+            input_prompt,
+            temperature=0.15,
+            top_p=0.1,
+            top_k=40,
+            repeat_penalty=1.1,
+            max_tokens=256,
+            stop=[
+                "<s>" "</s>" "\n" "[/INST]" "[INST]",
+                "### User:",
+                "### Assistant:",
+                "###",
+            ],
+            stream=True,
+        )
 
-        except Exception as e:
-            yield (
-                "🤖 El servidor no se encuentra disponible.",
-                "❌❌❌ Inténtalo de nuevo más tarde ❌❌❌",
-                "Error",
-            )
-            return (
-                "🤖 El servidor no se encuentra disponible.",
-                "❌❌❌ Inténtalo de nuevo más tarde ❌❌❌",
-                "Error",
-            )
+        temp = ""
+        for out in output:
+            stream = copy.deepcopy(out)
+            temp += stream["choices"][0]["text"]
+            yield title, temp, text
 
         cache_handler.add_to_cache(
             url=url, title=title, text=text, summary_type=mode, summary=temp
@@ -86,7 +154,7 @@ def generate_text(
 
 
 cache_handler = CacheHandler(max_cache_size=1000)
-feedback_log = gr.HuggingFaceDatasetSaver(auth_token, "Iker/Clickbait-News")
+hf_writer = gr.HuggingFaceDatasetSaver(auth_token, "Iker/Clickbait-News")
 
 demo = gr.Interface(
     generate_text,
@@ -141,7 +209,7 @@ Para obtener solo la respuesta al clickbait, selecciona 100""",
    
    🗒 La IA no es capaz de acceder a todas las webs, por ejemplo, si introduces un enlace a una noticia que requiere suscripción, la IA no podrá acceder a ella. Algunas webs pueden tener tecnologías para bloquear bots.
    
-   ⌚ La IA se encuentra corriendo en un hardware bastante modesto, debería tardar menos de 10 segundos en generar el resumen, pero si muchos usuarios usan la app a la vez, tendrás que esperar tu turno.
+   ⌚ La IA se encuentra corriendo en un hardware bastante modesto, por lo que puede tardar hasta un minuto en generar el resumen. Si muchos usuarios usan la app a la vez, tendrás que esperar tu turno.
    
    💸 Este es un projecto sin ánimo de lucro, no se genera ningún tipo de ingreso con esta app. Los datos, la IA y el código se publicarán para su uso en la investigación académica. No puedes usar esta app para ningún uso comercial.
    
@@ -151,7 +219,7 @@ Para obtener solo la respuesta al clickbait, selecciona 100""",
     concurrency_limit=1,
     allow_flagging="manual",
     flagging_options=[("👍", "correct"), ("👎", "incorrect")],
-    flagging_callback=feedback_log,
+    flagging_callback=hf_writer,
 )
 
 demo.queue(max_size=None)