Update app.py

app.py

@@ -9,14 +9,11 @@ You can also use efog 🌬️🌁🌫️SqlCoder by cloning this space. 🧬🔬
 Join us : 🌟TeamTonic🌟 is always making cool demos! Join our active builder's🛠️community 👻[![Let's build the future of AI together! 🚀🤖](https://discordapp.com/api/guilds/1109943800132010065/widget.png)](https://discord.gg/GWpVpekp) On 🤗Huggingface: [TeamTonic](https://huggingface.co/TeamTonic) & [MultiTransformer](https://huggingface.co/MultiTransformer) On 🌐Github: [Polytonic](https://github.com/tonic-ai) & contribute to 🌟 [Poly](https://github.com/tonic-ai/poly) 🤗Big thanks to Yuvi Sharma and all the folks at huggingface for the community grant 🤗
 """
 
-@spaces.GPU
-class SQLQueryGenerator:
-    def __init__(self, model_name, prompt_file="prompt.md", metadata_file="metadata.sql"):
-        self.tokenizer, self.model = self.get_tokenizer_model(model_name)
-        self.prompt_file = prompt_file
-        self.metadata_file = metadata_file
-    
-    def get_tokenizer_model(self, model_name):
+class TokenizerModel:
+    def __init__(self, model_name):
+        self.tokenizer, self.model = self.load_model(model_name)
+
+    def load_model(self, model_name):
         tokenizer = AutoTokenizer.from_pretrained(model_name)
         model = AutoModelForCausalLM.from_pretrained(
             model_name,
@@ -27,6 +24,12 @@ class SQLQueryGenerator:
         )
         return tokenizer, model
 
+class SQLQueryGenerator:
+    def __init__(self, tokenizer_model, prompt_file="prompt.md", metadata_file="metadata.sql"):
+        self.tokenizer_model = tokenizer_model
+        self.prompt_file = prompt_file
+        self.metadata_file = metadata_file
+
     def generate_prompt(self, question):
         with open(self.prompt_file, "r") as f:
             prompt = f.read()
@@ -39,14 +42,15 @@ class SQLQueryGenerator:
         )
         return prompt
 
+    @spaces.GPU
     def run_inference(self, question):
-        self.model.to('cuda')
+        self.tokenizer_model.model.to('cuda')
         prompt = self.generate_prompt(question)
-        eos_token_id = self.tokenizer.eos_token_id
+        eos_token_id = self.tokenizer_model.tokenizer.eos_token_id
         pipe = pipeline(
             "text-generation",
-            model=self.model,
-            tokenizer=self.tokenizer,
+            model=self.tokenizer_model.model,
+            tokenizer=self.tokenizer_model.tokenizer,
             max_new_tokens=300,
             do_sample=False,
             num_beams=5,
@@ -66,19 +70,17 @@ class SQLQueryGenerator:
         )
         return generated_query
 
-def generate_sql(question, sql_query_generator):
-    return sql_query_generator.run_inference(question)
-
 def main():
     model_name = "defog/sqlcoder2"
-    sql_query_generator = SQLQueryGenerator(model_name)
+    tokenizer_model = TokenizerModel(model_name)
+    sql_query_generator = SQLQueryGenerator(tokenizer_model)
 
     with gr.Blocks() as demo:
         gr.Markdown(title)
         question = gr.Textbox(label="Enter your question")
         submit = gr.Button("Generate SQL Query")
         output = gr.Textbox(label="🌬️🌁🌫️SqlCoder-2")
-        submit.click(fn=generate_sql, inputs=[question, gr.State(sql_query_generator)], outputs=output)
+        submit.click(fn=sql_query_generator.run_inference, inputs=question, outputs=output)
 
     demo.launch()