Update app.py

app.py

@@ -1,32 +1,70 @@
 import json
 
+import spaces
 import gradio as gr
+import torch
 from distilabel.llms import TransformersLLM
 from distilabel.steps.tasks.argillalabeller import ArgillaLabeller
+from transformers import BitsAndBytesConfig
 
-llm = TransformersLLM(model="microsoft/Phi-3-mini-4k-instruct")
+quantization_config = BitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_compute_dtype=torch.bfloat16,
+    bnb_4bit_use_double_quant=True,
+    bnb_4bit_quant_type="nf4",
+)
+
+
+llm = TransformersLLM(
+    model="microsoft/Phi-3-mini-4k-instruct",
+    torch_dtype="float16",
+    model_kwargs={
+        "quantization_config": quantization_config,
+        "device_map": "auto",
+    },
+)
 task = ArgillaLabeller(llm=llm)
 task.load()
 
+
+def load_examples():
+    with open("examples.json", "r") as f:
+        return json.load(f)
+
+
+# Create Gradio examples
+examples = load_examples()
+
+
+def process_fields(fields):
+    if isinstance(fields, str):
+        fields = json.loads(fields)
+    if isinstance(fields, dict):
+        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, field, question):
+def process_records_gradio(records, example_records, fields, question):
     try:
         # Convert string inputs to dictionaries
         records = json.loads(records)
         example_records = json.loads(example_records) if example_records else None
-        field = json.loads(field) if field else None
+        fields = process_fields(fields) if fields else None
         question = json.loads(question) if question else None
 
-        if not field and not question:
-            return "Error: Either field or question must be provided"
+        print(fields)
+        print(question)
+        print(example_records)
+
+        if not fields and not question:
+            return "Error: Either fields or question must be provided"
 
-        task.set_runtime_parameters(
-            {
-                "fields": [field] if field else None,
-                "question": question,
-                "example_records": example_records,
-            }
-        )
+        runtime_parameters = {"fields": fields, "question": question}
+        if example_records:
+            runtime_parameters["example_records"] = example_records
+        print(runtime_parameters)
+        task.set_runtime_parameters(runtime_parameters)
 
         results = []
         for record in records:
@@ -43,12 +81,13 @@ interface = gr.Interface(
     inputs=[
         gr.Code(label="Records (JSON)", language="json", lines=5),
         gr.Code(label="Example Records (JSON, optional)", language="json", lines=5),
-        gr.Code(label="Field (JSON, optional)", language="json"),
+        gr.Code(label="Fields (JSON, optional)", language="json"),
         gr.Code(label="Question (JSON, optional)", language="json"),
     ],
+    examples=examples,
     outputs=gr.Code(label="Suggestions", language="json", lines=10),
     title="Record Processing Interface",
-    description="Enter JSON data for records, example records, field, and question. At least one of field or question must be provided.",
+    description="Enter JSON data for `rg.Record.to_dict()`, `List[rg.Record.to_dict()]`, `List[Field].serialize()`, or `List[rg.Question.serialize()]` At least one of fields or question must be provided.",
 )
 
 if __name__ == "__main__":