Team Ai
Apppublic

daljeetsingh/SQL_Generation

sourceHugging Facemitupdated 2y agoView on Hugging Face
3likes
app.py76 linesDownload Raw Back to root
1## https://www.kaggle.com/code/unravel/fine-tuning-of-a-sql-model2 3import spaces4from transformers import AutoModelForSeq2SeqLM, AutoTokenizer5import gradio as gr6import torch7from transformers.utils import logging8from example_queries import small_query, long_query9 10logging.set_verbosity_info()11logger = logging.get_logger("transformers")12 13model_name='t5-small'14tokenizer = AutoTokenizer.from_pretrained(model_name)15original_model = AutoModelForSeq2SeqLM.from_pretrained(model_name, torch_dtype=torch.bfloat16)16 17ft_model_name="daljeetsingh/sql_ft_t5small_kag" #"cssupport/t5-small-awesome-text-to-sql"18ft_model = AutoModelForSeq2SeqLM.from_pretrained(ft_model_name, torch_dtype=torch.bfloat16)19 20original_model.to('cuda')21ft_model.to('cuda')22 23@spaces.GPU24def translate_text(text):25   prompt = f"{text}"26   inputs = tokenizer(prompt, return_tensors='pt')27   inputs = inputs.to('cuda')28 29   try:30        output = tokenizer.decode(31            original_model.generate(32                inputs["input_ids"], 33                max_new_tokens=200,34            )[0], 35            skip_special_tokens=True36        )37        ft_output = tokenizer.decode(38            ft_model.generate(39                inputs["input_ids"], 40                max_new_tokens=200,41            )[0], 42            skip_special_tokens=True43        )44        return [output, ft_output]45   except Exception as e:46       return f"Error: {str(e)}"47 48 49with gr.Blocks() as demo:50    with gr.Row():51        with gr.Column():52            prompt = gr.Textbox(53                value=small_query,54                lines=8,55                placeholder="Enter prompt...",56                label="Prompt"57            )58            submit_btn = gr.Button(value="Generate")59        with gr.Column():60            orig_output = gr.Textbox(label="OriginalModel", lines=2) 61            ft_output = gr.Textbox(label="FTModel", lines=8)62 63    submit_btn.click(64        translate_text, inputs=[prompt], outputs=[orig_output, ft_output], api_name=False65    )66    examples = gr.Examples(67        examples=[68            [small_query],69            [long_query],70        ],71        inputs=[prompt],72    )73 74demo.launch(show_api=False, share=True, debug=True)75 76