| import replicate | |
| class Text2SQLModel: | |
| """ | |
| A class representing a Text-to-SQL model for generating SQL queries from LLM. | |
| """ | |
| def __init__(self): | |
| pass | |
| def load_model(self): | |
| """Loads the machine learning model for Text-to-SQL processing.""" | |
| pass | |
| def generate_query(self, prompt): | |
| output = replicate.run( | |
| "ns-dev-sentience/sqlcoder:18bcabd866a64547daf3c6044cdebbd47a1f489571110087b80722848eb09398", | |
| input={"prompt": prompt}, | |
| ) | |
| return ( | |
| output.split("```PostgresSQL")[-1].split("```")[0].split(";")[0].strip() | |
| + ";" | |
| ) | |