daljeetsingh/SQL_Generation
3
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 