Upload 4 files

nlp/sql_model.py

@@ -1,14 +1,14 @@
-# from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
-import os
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+# import os
+# from transformers import AutoTokenizer, AutoModelForCausalLM
 
 class SQLModel:
     def __init__(self, model_name="google/flan-t5-base"):
-        # self.tokenizer = AutoTokenizer.from_pretrained(model_name)
-        # self.model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
-        hf_token = os.environ.get("HF_HUB_TOKEN")
-        self.tokenizer = AutoTokenizer.from_pretrained("google/gemma-2b-it",use_auth_token=hf_token)
-        self.model = AutoModelForCausalLM.from_pretrained("google/gemma-2b-it",use_auth_token=hf_token)
+        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
+        self.model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
+        # hf_token = os.environ.get("HF_HUB_TOKEN")
+        # self.tokenizer = AutoTokenizer.from_pretrained("google/gemma-2b-it",use_auth_token=hf_token)
+        # self.model = AutoModelForCausalLM.from_pretrained("google/gemma-2b-it",use_auth_token=hf_token)
 
     def generate_sql(self, natural_language_query):
         input_text = f"""You are a highly skilled SQL translator. Your task is to convert natural language descriptions of data queries into correct and optimized SQL statements.