| import streamlit as st |
| from llama_cpp import Llama |
| from sql import get_table_schema |
|
|
|
|
| @st.cache_resource() |
| def load_llm(repo_id, filename): |
| llm = Llama.from_pretrained( |
| repo_id=repo_id, |
| filename=filename, |
| verbose=True, |
| use_mmap=True, |
| use_mlock=True, |
| n_threads=4, |
| n_threads_batch=4, |
| n_ctx=8000, |
| ) |
| print(f"{repo_id} loaded successfully. ✅") |
| return llm |
|
|
|
|
| def generate_system_prompt(table_name, table_schema): |
| """ |
| Generates a prompt to provide context about a table's schema for LLM to convert natural language to SQL. |
| |
| Args: |
| table_name (str): The name of the table. |
| table_schema (list): A list of tuples where each tuple contains information about the columns in the table. |
| |
| Returns: |
| str: The generated prompt to be used by the LLM. |
| """ |
| prompt = f"""You are an expert in writing SQL queries for relational databases, specially sqlite. |
| You will be provided with a database schema and a natural |
| language question, and your task is to generate an accurate SQL query. |
| |
| The database has a table named '{table_name}' with the following schema:\n\n""" |
|
|
| prompt += "Columns:\n" |
|
|
| for col in table_schema: |
| column_name = col[1] |
| column_type = col[2] |
| prompt += f"- {column_name} ({column_type})\n" |
|
|
| prompt += "\nGenerate a SQL query based on the following natural language question. ONLY return the SQL query and nothing else." |
|
|
| return prompt |
|
|
|
|
| |
| def response_generator(llm, messages, question, table_name, db_name): |
| table_schema = get_table_schema(db_name, table_name) |
| llm_prompt = generate_system_prompt(table_name, table_schema) |
| user_prompt = f"""Question: {question}""" |
|
|
| print(messages, llm_prompt, user_prompt) |
| history = [{"content": llm_prompt.format(table_name=table_name), "role": "system"}] |
|
|
| for val in messages: |
| history.append(val) |
|
|
| history.append({"role": "user", "content": user_prompt}) |
|
|
| response = llm.create_chat_completion( |
| messages=history, |
| max_tokens=2048, |
| temperature=0.7, |
| top_p=0.95, |
| ) |
| answer = response["choices"][0]["message"]["content"] |
|
|
| query = answer.replace("```sql", "").replace("```", "") |
| query = query.strip() |
| return query |
|
|
