feat: use InferenceEndpointsLLM

app.py

@@ -1,36 +1,13 @@
 import json
-import os
 
 import gradio as gr
-from distilabel.llms import LlamaCppLLM
+from distilabel.llms import InferenceEndpointsLLM
 from distilabel.steps.tasks.argillalabeller import ArgillaLabeller
 
-file_path = os.path.join(os.path.dirname(__file__), "Qwen2-5-0.5B-Instruct-f16.gguf")
-download_url = "https://huggingface.co/gaianet/Qwen2.5-0.5B-Instruct-GGUF/resolve/main/Qwen2.5-0.5B-Instruct-Q8_0.gguf?download=true"
-
-
-if not os.path.exists(file_path):
-    import requests
-    import tqdm
-
-    response = requests.get(download_url, stream=True)
-    total_length = int(response.headers.get("content-length"))
-
-    with open(file_path, "wb") as f:
-        for chunk in tqdm.tqdm(
-            response.iter_content(chunk_size=1024 * 1024),
-            total=total_length / (1024 * 1024),
-            unit="KB",
-            unit_scale=True,
-        ):
-            f.write(chunk)
-
-context_window = 1024 * 128
-llm = LlamaCppLLM(
-    model_path=file_path,
-    n_gpu_layers=-1,
-    n_ctx=context_window,
-    generation_kwargs={"max_new_tokens": context_window},
+llm = InferenceEndpointsLLM(
+    model_id="meta-llama/Meta-Llama-3.1-8B-Instruct",
+    tokenizer_id="meta-llama/Meta-Llama-3.1-8B-Instruct",
+    generation_kwargs={"max_new_tokens": 1000 * 128},
 )
 task = ArgillaLabeller(llm=llm)
 task.load()