Upload 4 files

nlp/sql_model.py

@@ -1,12 +1,14 @@
 # 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)
-        self.tokenizer = AutoTokenizer.from_pretrained("google/gemma-7b-it")
-        self.model = AutoModelForCausalLM.from_pretrained("google/gemma-7b-it",)
+        hf_token = os.environ.get("HF_HUB_TOKEN")
+        self.tokenizer = AutoTokenizer.from_pretrained("google/gemma-7b-it",use_auth_token=hf_token)
+        self.model = AutoModelForCausalLM.from_pretrained("google/gemma-7b-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.