|
import gradio as gr |
|
from transformers import AutoModelForCausalLM, AutoTokenizer |
|
|
|
|
|
chatbot_model_name = "facebook/bart-large-mnli" |
|
chatbot_tokenizer = AutoTokenizer.from_pretrained(chatbot_model_name) |
|
chatbot_model = AutoModelForCausalLM.from_pretrained(chatbot_model_name) |
|
|
|
|
|
sql_model_name = "your_sql_model_name" |
|
sql_tokenizer = AutoTokenizer.from_pretrained(sql_model_name) |
|
sql_model = AutoModelForCausalLM.from_pretrained(sql_model_name) |
|
|
|
def chatbot_response(user_message): |
|
|
|
inputs = chatbot_tokenizer.encode("User: " + user_message, return_tensors="pt") |
|
outputs = chatbot_model.generate(inputs, max_length=100, num_return_sequences=1) |
|
response = chatbot_tokenizer.decode(outputs[0], skip_special_tokens=True) |
|
|
|
return response |
|
|
|
def execute_sql(user_query): |
|
|
|
inputs = sql_tokenizer(user_query, return_tensors="pt") |
|
outputs = sql_model.generate(inputs['input_ids'], attention_mask=inputs['attention_mask'], max_length=1000) |
|
response = sql_tokenizer.decode(outputs[0], skip_special_tokens=True) |
|
|
|
return response |
|
|
|
|
|
chatbot_interface = gr.Interface( |
|
fn=chatbot_response, |
|
inputs=gr.Textbox(prompt="You:"), |
|
outputs=gr.Textbox(), |
|
live=True, |
|
capture_session=True, |
|
title="Chatbot", |
|
description="Type your message in the box above, and the chatbot will respond.", |
|
) |
|
|
|
sql_execution_interface = gr.Interface( |
|
fn=execute_sql, |
|
inputs=gr.Textbox(prompt="Enter your SQL query:"), |
|
outputs=gr.Textbox(), |
|
live=True, |
|
capture_session=True, |
|
title="SQL Execution", |
|
description="Type your SQL query in the box above, and the chatbot will execute it.", |
|
) |
|
|
|
|
|
combined_interface = gr.Interface([chatbot_interface, sql_execution_interface], layout="horizontal") |
|
|
|
|
|
if __name__ == "__main__": |
|
combined_interface.launch() |
|
|