Description
This model can generate SQL queries from natural questions and custom database schemas with a single table. It is based on a large-size LLM, which is finetuned by John Snow Labs on a dataset having schemas with single tables.
Predicted Entities
How to use
question = "Calculate the average age of patients with blood type 'A-' "
query_schema = {
"patient": ["ID","Name","Age","Gender","BloodType","Weight","Height","Address","Email","Phone"]
}
document_assembler = DocumentAssembler()\
.setInputCol("text")\
.setOutputCol("document")
text2sql_with_schema_single_table = Text2SQL.pretrained("text2sql_with_schema_single_table", "en", "clinical/models")\
.setMaxNewTokens(200)\
.setSchema(query_schema)\
.setInputCols(["document"])\
.setOutputCol("sql")
pipeline = Pipeline(stages=[document_assembler, text2sql_with_schema_single_table])
data = spark.createDataFrame([[question]]).toDF("text")
pipeline.fit(data)\
.transform(data)\
.select("sql.result")\
.show(truncate=False)
val question = """Calculate the average age of patients with blood type 'A-' """
val query_schema : Map[String, List[String]] = Map(
"patient" -> List("ID", "Name", "Age", "Gender", "BloodType", "Weight", "Height", "Address", "Email", "Phone")
)
val document_assembler = new DocumentAssembler()
.setInputCol("text")
.setOutputCol("document")
val text2sql_with_schema_single_table = new Text2SQL.pretrained("text2sql_with_schema_single_table", "en", "clinical/models")
.setMaxNewTokens(200)
.setSchema(query_schema)
.setInputCols(["document"])
.setOutputCol("sql")
val pipeline = new Pipeline().setStages(Array(document_assembler, text2sql_with_schema_single_table ))
val data = Seq(Array(text)).toDS.toDF("text")
val result = pipeline.fit(data).transform(data)
Results
[SELECT AVG(Age) FROM patient WHERE BloodType = "A-"]
Model Information
| Model Name: | text2sql_with_schema_single_table |
| Compatibility: | Healthcare NLP 5.0.1+ |
| License: | Licensed |
| Edition: | Official |
| Input Labels: | [document_question] |
| Output Labels: | [sql_query] |
| Language: | en |
| Size: | 3.0 GB |