File size: 5,346 Bytes
39860ea aae3c09 39860ea aae3c09 39860ea aae3c09 39860ea aae3c09 ac038b5 39860ea aae3c09 39860ea aae3c09 39860ea aae3c09 39860ea aae3c09 39860ea aae3c09 39860ea aae3c09 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 |
import json
import os
import torch
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
# Set parameters
max_seq_length = 4096
dtype = None
load_in_4bit = False
# Define the base model identifier (the full model from Hugging Face Hub)
base_model_id = "meta-llama/Llama-3.2-3B-Instruct"
# 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,
)
# 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()
# 3. Prepare the merged model for inference.
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 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})
|