Update app.py

app.py

@@ -3,12 +3,34 @@ import json
 import spaces
 import gradio as gr
 import torch
-from distilabel.llms import vLLM
+from distilabel.llms import LlamaCppLLM
 from distilabel.steps.tasks.argillalabeller import ArgillaLabeller
 
+file_path = os.path.join(os.path.dirname(__file__), "qwen2-0_5b-instruct-fp16.gguf")
+download_url = "https://huggingface.co/Qwen/Qwen2-0.5B-Instruct-GGUF/resolve/main/qwen2-0_5b-instruct-fp16.gguf?download=true"
 
-llm = vLLM(
-    model="Qwen/Qwen2-0.5B-Instruct-GGUF"
+
+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),
+            total=total_length,
+            unit="KB",
+            unit_scale=True,
+        ):
+            f.write(chunk)
+
+
+llm = LlamaCppLLM(
+    model_path=file_path,
+    n_gpu_layers=-1,
+    n_ctx=1024 * 4,
 )
 task = ArgillaLabeller(llm=llm)
 task.load()
@@ -30,7 +52,7 @@ def process_fields(fields):
         fields = [fields]
     return [field if isinstance(field, dict) else json.loads(field) for field in fields]
 
-
+@spaces.GPU
 def process_records_gradio(records, example_records, fields, question):
     try:
         # Convert string inputs to dictionaries