Upload 2 files

inference.py

@@ -0,0 +1,121 @@
+import json
+from unsloth import FastLanguageModel
+from transformers import TextStreamer  # if needed elsewhere
+
+# Set parameters
+max_seq_length = 4096 
+dtype = None 
+load_in_4bit = False 
+
+# Load model and tokenizer once at startup
+model, tokenizer = FastLanguageModel.from_pretrained(
+    model_name="betterdataai/large-tabular-model", 
+    max_seq_length=max_seq_length,
+    dtype=dtype,
+    load_in_4bit=load_in_4bit,
+)
+FastLanguageModel.for_inference(model)
+
+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)
+
+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)
+
+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 desired 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})

requirements.txt

@@ -0,0 +1,9 @@
+unsloth 
+pandas
+datasets
+trl
+scipy
+transformers>=4.30.0
+peft>=0.4.0
+accelerate>=0.20.0
+torch>=2.0