Upload handler.py

handler.py

@@ -1,86 +1,138 @@
-import os
 import json
+import os
 import torch
-from peft import PeftModel
-from transformers import (
-    LlamaForCausalLM,
-    LlamaTokenizer,
-    GenerationConfig,
-)
+from transformers import TextStreamer  # if needed elsewhere
+from unsloth import FastLanguageModel  # Assumes FastLanguageModel supports loading a base model
+from peft import PeftModel  # For loading the adapter onto the base model
 
-class EndpointHandler:
-    def __init__(self, model_dir: str = ".", **kwargs):
-        """
-        This method runs once when the Endpoint first starts.
-        - model_dir is the local directory of *this* repository
-          which contains your LoRA adapter weights (e.g. adapter_model.safetensors).
-        """
+# Set parameters
+max_seq_length = 4096 
+dtype = None 
+load_in_4bit = False 
 
-        # 1) Base model from Hugging Face
-        #    Make sure to use the EXACT base you trained on, or it won't match your LoRA.
-        self.base_model_id = "unsloth/Llama-3.2-3B-Instruct"
-        
-        # If your base model is gated/private, you'll need a token:
-        # hf_token = os.getenv("HF_TOKEN", None)
-        
-        # 2) Load the tokenizer
-        self.tokenizer = LlamaTokenizer.from_pretrained(
-            self.base_model_id,
-            trust_remote_code=True,
-            # use_auth_token=hf_token,  # if needed
-        )
-        
-        # 3) Load the base model
-        self.base_model = LlamaForCausalLM.from_pretrained(
-            self.base_model_id,
-            device_map="auto",         # or "cuda:0"
-            torch_dtype=torch.float16, # or bfloat16
-            trust_remote_code=True,
-            # use_auth_token=hf_token, # if needed
-        )
+# Define the base model identifier (the full model from Hugging Face Hub)
+base_model_id = "meta-llama/Llama-3.2-3B-Instruct"
 
-        # 4) Load/merge your LoRA adapter
-        self.model = PeftModel.from_pretrained(
-            self.base_model,
-            model_dir,                 # The local directory of this repo
-            torch_dtype=torch.float16,
-        ).eval()
-
-    def __call__(self, data):
-        """
-        This method is called for every request to the endpoint.
-        `data` is a dictionary (or JSON string) containing user inputs.
-        Returns a dictionary or string (will be serialized as JSON).
-        """
-        # If data is a JSON string, parse it:
-        if isinstance(data, str):
-            data = json.loads(data)
-
-        # Extract the user prompt from the request payload
-        prompt = data.get("inputs", "")
-        if not isinstance(prompt, str):
-            raise ValueError("`inputs` must be a string.")
+# 1. Load the base model and tokenizer from the Hub.
+# (This downloads the complete base model with all weights.)
+base_model, tokenizer = FastLanguageModel.from_pretrained(
+    model_name=base_model_id, 
+    max_seq_length=max_seq_length,
+    dtype=dtype,
+    load_in_4bit=load_in_4bit,
+)
 
-        # Optionally extract generation params (max_new_tokens, temperature, etc.)
-        # If none provided, use defaults:
-        gen_params = data.get("parameters", {})
-        generation_config = GenerationConfig(
-            max_new_tokens=gen_params.get("max_new_tokens", 128),
-            temperature=gen_params.get("temperature", 0.7),
-            top_p=gen_params.get("top_p", 0.9),
-            do_sample=gen_params.get("do_sample", True),
-            # etc.
-        )
+# 2. Load your LoRA adapter weights from your repository.
+#    Here, "betterdataai/large-tabular-model" should be the local directory or identifier where the adapter weights reside.
+#    Ensure that this path contains the adapter weights (e.g. adapter_model.safetensors) and configuration.
+model = PeftModel.from_pretrained(
+    base_model, 
+    "betterdataai/large-tabular-model",  # Path to your adapter weights
+    torch_dtype=torch.float16,
+).eval()
 
-        # Tokenize the prompt
-        inputs = self.tokenizer(prompt, return_tensors="pt").to(self.model.device)
+# 3. Prepare the merged model for inference.
+FastLanguageModel.for_inference(model)
 
-        # Generate text
-        with torch.no_grad():
-            output_ids = self.model.generate(**inputs, generation_config=generation_config)
+def prompt_transformation(prompt):
+    initial_prompt = """
+    We have the following natural language query:
+    "{}"
+    
+    Transform the above natural language query into a formalized prompt format. The format should include:
+    
+    1. A sentence summarizing the objective.
+    2. A description of the columns, including their data types and examples.
+    3. Four example rows of the dataset in CSV format.
+    
+    An example of this format is as follows, please only focus on the format, not the content:
+    
+    "You are tasked with generating a synthetic dataset based on the following description. The dataset represents employee information. The dataset should include the following columns:
+    
+    - NAME (String): Employee's full name, consisting of a first and last name (e.g., "John Doe", "Maria Lee", "Wei Zhang").
+    - GENDER (String): Employee's gender (e.g., "Male", "Female").
+    - EMAIL (String): Employee's email address, following the standard format.
+    - CITY (String): City where the employee resides (e.g., "New York", "London", "Beijing").
+    - COUNTRY (String): Country where the employee resides (e.g., "USA", "UK", "China").
+    - SALARY (Float): Employee's annual salary, a value between 30000 and 150000 (e.g., 55000.0, 75000.0).
+    
+    Here are some examples:
+    NAME,GENDER,EMAIL,CITY,COUNTRY,SALARY
+    John Doe,Male,john.doe@example.com,New York,USA,56000.0
+    Maria Lee,Female,maria.lee@nus.edu.sg,London,UK,72000.0
+    Wei Zhang,Male,wei.zhang@meta.com,Beijing,China,65000.0
+    Sara Smith,Female,sara.smith@orange.fr,Paris,France,85000.0"
+    
+    Here is the transformed query from the given natural language query:
+    """
+    messages = [
+        {"role": "system", "content": initial_prompt.format(prompt)},
+        {"role": "user", "content": "transform the given natural language text to the designated format"}
+    ]
+    
+    inputs = tokenizer.apply_chat_template(
+        messages,
+        tokenize=True,
+        add_generation_prompt=True,  # Required for generation
+        return_tensors="pt",
+    ).to("cuda")
+    
+    output_ids = model.generate(
+        input_ids=inputs, 
+        max_new_tokens=4096, 
+        use_cache=True, 
+        temperature=1.5, 
+        min_p=0.1
+    )
+    
+    generated_ids = output_ids[0][inputs.shape[1]:]
+    return tokenizer.decode(generated_ids, skip_special_tokens=True)
 
-        # Decode the output
-        output_text = self.tokenizer.decode(output_ids[0], skip_special_tokens=True)
+def table_generation(prompt):
+    messages = [
+        {"role": "system", "content": prompt},
+        {"role": "user", "content": "create 20 data rows"}
+    ]
+    
+    inputs = tokenizer.apply_chat_template(
+        messages,
+        tokenize=True,
+        add_generation_prompt=True,  # Required for generation
+        return_tensors="pt",
+    ).to("cuda")
+    
+    output_ids = model.generate(
+        input_ids=inputs, 
+        max_new_tokens=4096, 
+        use_cache=True, 
+        temperature=1.5, 
+        min_p=0.1
+    )
+    
+    generated_ids = output_ids[0][inputs.shape[1]:]
+    return tokenizer.decode(generated_ids, skip_special_tokens=True)
 
-        # Return the generated text in a JSON-friendly format
-        return {"generated_text": output_text}
+def predict(input_data):
+    """
+    Inference endpoint entry point.
+    
+    Expects input_data as a JSON string or dict with a key "query" that contains the natural language query.
+    Returns a JSON string with the generated table.
+    """
+    try:
+        if isinstance(input_data, str):
+            data = json.loads(input_data)
+        else:
+            data = input_data
+        user_query = data.get("query", "")
+    except Exception:
+        return json.dumps({
+            "error": "Invalid input format. Please provide a JSON payload with a 'query' field."
+        })
+    
+    # Transform the user query into the designated prompt format.
+    transformed_prompt = prompt_transformation(user_query)
+    # Generate the table using the transformed prompt.
+    generated_table = table_generation(transformed_prompt)
+    
+    return json.dumps({"result": generated_table})