Update app.py

app.py

@@ -9,78 +9,67 @@ 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 🤗
 """
 
-class TokenizerModel:
-    def __init__(self, model_name):
-        self.tokenizer, self.model = self.load_model(model_name)
+def load_tokenizer_model(model_name):
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    model = AutoModelForCausalLM.from_pretrained(
+        model_name,
+        trust_remote_code=True,
+        torch_dtype=torch.float16,
+        device_map="auto",
+        use_cache=True,
+    )
+    return tokenizer, model
 
-    def load_model(self, model_name):
-        tokenizer = AutoTokenizer.from_pretrained(model_name)
-        model = AutoModelForCausalLM.from_pretrained(
-            model_name,
-            trust_remote_code=True,
-            torch_dtype=torch.float16,
-            device_map="auto",
-            use_cache=True,
-        )
-        return tokenizer, model
+def generate_prompt(question, prompt_file="prompt.md", metadata_file="metadata.sql"):
+    with open(prompt_file, "r") as f:
+        prompt = f.read()
 
-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
+    with open(metadata_file, "r") as f:
+        table_metadata_string = f.read()
 
-    def generate_prompt(self, question):
-        with open(self.prompt_file, "r") as f:
-            prompt = f.read()
+    prompt = prompt.format(
+        user_question=question, table_metadata_string=table_metadata_string
+    )
+    return prompt
 
-        with open(self.metadata_file, "r") as f:
-            table_metadata_string = f.read()
-
-        prompt = prompt.format(
-            user_question=question, table_metadata_string=table_metadata_string
-        )
-        return prompt
-
-    @spaces.GPU
-    def run_inference(self, question):
-        self.tokenizer_model.model.to('cuda')
-        prompt = self.generate_prompt(question)
-        eos_token_id = self.tokenizer_model.tokenizer.eos_token_id
-        pipe = pipeline(
-            "text-generation",
-            model=self.tokenizer_model.model,
-            tokenizer=self.tokenizer_model.tokenizer,
-            max_new_tokens=300,
-            do_sample=False,
-            num_beams=5,
-        )
-        generated_query = (
-            pipe(
-                prompt,
-                num_return_sequences=1,
-                eos_token_id=eos_token_id,
-                pad_token_id=eos_token_id,
-            )[0]["generated_text"]
-            .split("```sql")[-1]
-            .split("```")[0]
-            .split(";")[0]
-            .strip()
-            + ";"
-        )
-        return generated_query
+@spaces.GPU
+def run_inference(question, model, tokenizer):
+    model.to('cuda')
+    prompt = generate_prompt(question)
+    eos_token_id = tokenizer.eos_token_id
+    pipe = pipeline(
+        "text-generation",
+        model=model,
+        tokenizer=tokenizer,
+        max_new_tokens=300,
+        do_sample=False,
+        num_beams=5,
+    )
+    generated_query = (
+        pipe(
+            prompt,
+            num_return_sequences=1,
+            eos_token_id=eos_token_id,
+            pad_token_id=eos_token_id,
+        )[0]["generated_text"]
+        .split("```sql")[-1]
+        .split("```")[0]
+        .split(";")[0]
+        .strip()
+        + ";"
+    )
+    return generated_query
 
 def main():
     model_name = "defog/sqlcoder2"
-    tokenizer_model = TokenizerModel(model_name)
-    sql_query_generator = SQLQueryGenerator(tokenizer_model)
+    tokenizer, model = load_tokenizer_model(model_name)
 
     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=sql_query_generator.run_inference, inputs=question, outputs=output)
+        submit.click(fn=lambda x: run_inference(x, model, tokenizer), inputs=question, outputs=output)
 
     demo.launch()