Upload 2 files

Dockerfile

@@ -10,4 +10,4 @@ ENV HOME = /home/user \
 WORKDIR $HOME/app
 COPY --chown==user . $HOME/app
 
-CMD [ "uvicorn", "app:app", "--host", "0.0.0.0", "--port", "7860"  ]
+CMD [ "uvicorn", "main:app", "--host", "0.0.0.0", "--port", "7860"  ]

main.py

@@ -0,0 +1,70 @@
+from fastapi import FastAPI
+from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
+from typing import List, Dict
+import time
+import datetime
+import uvicorn
+
+model = AutoModelForSeq2SeqLM.from_pretrained("KN123/nl2sql")
+tokenizer = AutoTokenizer.from_pretrained("KN123/nl2sql")
+
+def get_prompt(tables, question):
+    prompt = f"""convert question and table into SQL query. tables: {tables}. question: {question}"""
+    # print(prompt)
+    return prompt
+
+def prepare_input(question: str, tables: Dict[str, List[str]]):
+    tables = [f"""{table_name}({",".join(tables[table_name])})""" for table_name in tables]
+    # print(tables)
+    tables = ", ".join(tables)
+    # print(tables)
+    prompt = get_prompt(tables, question)
+    # print(prompt)
+    input_ids = tokenizer(prompt, max_length=512, return_tensors="pt").input_ids
+    # print(input_ids)
+    return input_ids
+
+def inference(question: str, tables: Dict[str, List[str]]) -> str:
+    input_data = prepare_input(question=question, tables=tables)
+    input_data = input_data.to(model.device)
+    outputs = model.generate(inputs=input_data, num_beams=10, top_k=10, max_length=512)
+    # print("Outputs", outputs)
+    result = tokenizer.decode(token_ids=outputs[0], skip_special_tokens=True)
+    return result
+
+app = FastAPI()
+
+@app.get("/")
+def home():
+    return {
+        "message" : "Hello there! Everything is working fine!",
+        "api-version": "1.0.0",
+        "role": "nl2sql",
+        "description": "This api can be used to convert natural language to SQL given the human prompt, tables and the attributes."
+    }
+
+@app.get("/generate")
+def generate(text:str):
+    start = time.time()
+    res = inference("how many people with name jui and age less than 25", {
+        "people_name":["id","name"], "people_age": ["people_id","age"]
+    })
+    end = time.time()
+    total_time_taken = end - start
+    current_utc_datetime = datetime.datetime.now(datetime.timezone.utc)
+    current_date = datetime.date.today()
+    timezone_name = time.tzname[time.daylight]
+    print(res)
+    return {
+        "api_response": f"{res}",
+        "time_taken(s)": f"{total_time_taken}",
+        "request_details": {
+            "utc_datetime": f"{current_utc_datetime}",
+            "current_date": f"{current_date}",
+            "timezone_name": f"{timezone_name}"
+        }
+    }
+
+
+if __name__ == "__main__":
+    uvicorn.run(app, host="127.0.0.1", port=8000)