Upload app.py

app.py

@@ -1,22 +1,62 @@
 import os
-
+import requests
 import gradio as gr
 from download_url import download_text_and_title
 from cache_system import CacheHandler
-from gradio_client import Client
 from collections import OrderedDict
-from typing import Any
+from typing import Any, Iterable, List
 import datetime
+import json
 
-
-server = os.environ.get("SERVER") or True
+server = os.environ.get("SERVER") or "http://localhost:7861/generate"
 auth_token = os.environ.get("TOKEN") or True
+API_KEY = os.environ.get("API_KEY") or None
 
-client = Client(server, hf_token=auth_token, verbose=False)
 
 total_runs = 0
 
 
+def call_vllm_server(tittle, body, mode, stream=True):
+    api_url = "http://localhost:7861/generate"
+    headers = {"User-Agent": "Test Client"}
+
+    json = {
+        "n": 1,
+        "tittle": tittle,
+        "body": body,
+        "mode": mode,
+        "max_tokens": 4096,
+        "temperature": 0.15,
+        "top_p": 0.1,
+        "top_k": 40,
+        "repetition_penalty": 1.1,
+        "stop": [
+            "<s>",
+            "</s>",
+            "\n",
+            "[/INST]",
+            "[INST]",
+            "### User:",
+            "### Assistant:",
+            "###",
+        ],
+        "stream": stream,
+        "api_key": API_KEY,
+    }
+    response = requests.post(api_url, headers=headers, json=json)
+    return response
+
+
+def get_streaming_response(response: requests.Response) -> Iterable[List[str]]:
+    for chunk in response.iter_lines(
+        chunk_size=8192, decode_unicode=False, delimiter=b"\0"
+    ):
+        if chunk:
+            data = json.loads(chunk.decode("utf-8"))
+            output = data["text"]
+            yield output
+
+
 class HuggingFaceDatasetSaver_custom(gr.HuggingFaceDatasetSaver):
     def _deserialize_components(
         self,
@@ -88,6 +128,7 @@ def generate_text(
         try:
             title, text, url = download_text_and_title(url)
         except Exception as e:
+            print(e)
             title = None
             text = None
 
@@ -113,18 +154,13 @@ def generate_text(
         progress(0.5, desc="🤖 Leyendo noticia")
 
         try:
-            temp = client.submit(
-                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",
-            )
-
-            for o in temp:
-                yield title, o, text
+            response = call_vllm_server(title, text, mode, stream=True)
+            for h in get_streaming_response(response):
+                temp = h[0]
+                yield title, temp, text
 
         except Exception as e:
+            print(e)
             yield (
                 "🤖 El servidor no se encuentra disponible.",
                 "❌❌❌ Inténtalo de nuevo más tarde ❌❌❌",
@@ -136,7 +172,6 @@ def generate_text(
                 "Error",
             )
 
-        temp = temp.outputs()[-1]
         cache_handler.add_to_cache(
             url=url, title=title, text=text, summary_type=mode, summary=temp
         )
@@ -200,7 +235,7 @@ Para obtener solo la respuesta al clickbait, selecciona 100""",
             placeholder="Aquí aparecerá el resumen de la noticia.",
         ),
     ],
-    #title="⚔️ Clickbait Fighter! ⚔️",
+    # title="⚔️ Clickbait Fighter! ⚔️",
     thumbnail="https://huggingface.co/spaces/Iker/ClickbaitFighter/resolve/main/logo2.png",
     theme="JohnSmith9982/small_and_pretty",
     description="""
@@ -220,11 +255,10 @@ Para obtener solo la respuesta al clickbait, selecciona 100""",
    🧪 El modelo se encuentra en fase de desarrollo, si quieres ayudar a mejorarlo puedes usar los botones 👍 y 👎 para valorar el resumen. ¡Gracias por tu ayuda!""",
     article="Esta Inteligencia Artificial ha sido generada por Iker García-Ferrero. Puedes saber más sobre mi trabajo en mi [página web](https://ikergarcia1996.github.io/Iker-Garcia-Ferrero/) o mi perfil de [X](https://twitter.com/iker_garciaf). Puedes ponerte en contacto conmigo a través de correo electrónico (ver web) y X.",
     cache_examples=False,
-    concurrency_limit=1,
     allow_flagging="manual",
     flagging_options=[("👍", "correct"), ("👎", "incorrect")],
     flagging_callback=hf_writer,
 )
 
-demo.queue(max_size=None)
+demo.queue(max_size=None, concurrency_count=100)
 demo.launch(share=False)