feat add model prediction for text2sql prompt

prediction.py

@@ -18,7 +18,7 @@ else:
 
 
 from transformers import pipeline as hf_pipeline
-import torch
+
 import litellm
 
 from tqdm import tqdm
@@ -27,18 +27,25 @@ from tqdm import tqdm
 class ModelPrediction:
     def __init__(self):
         self.model_name2pred_func = {
-            "gpt-3.5": self._model_prediction("gpt-3.5"),
-            "gpt-4o-mini": self._model_prediction("gpt-4o-mini"),
-            "o1-mini": self._model_prediction("o1-mini"),
-            "QwQ": self._model_prediction("QwQ"),
-            "DeepSeek-R1-Distill-Llama-70B": self._model_prediction(
+            "gpt-3.5": self._init_model_prediction("gpt-3.5"),
+            "gpt-4o-mini": self._init_model_prediction("gpt-4o-mini"),
+            "o1-mini": self._init_model_prediction("o1-mini"),
+            "QwQ": self._init_model_prediction("QwQ"),
+            "DeepSeek-R1-Distill-Llama-70B": self._init_model_prediction(
                 "DeepSeek-R1-Distill-Llama-70B"
             ),
-            "llama-8": self._model_prediction("llama-8"),
+            "llama-8": self._init_model_prediction("llama-8"),
         }
 
         self._model_name = None
         self._pipeline = None
+        self.base_prompt= (
+            "Translate the following question in SQL code to be executed over the database to fetch the answer. Return the sql code in ```sql ```\n"
+            " Question\n"
+            "{question}\n"
+            "Database Schema\n"
+            "{db_schema}\n"
+        )
 
     @property
     def pipeline(self):
@@ -46,7 +53,6 @@ class ModelPrediction:
             self._pipeline = hf_pipeline(
                 task="text-generation",
                 model=self._model_name,
-                torch_dtype=torch.bfloat16,
                 device_map="auto",
             )
         return self._pipeline
@@ -67,14 +73,8 @@ class ModelPrediction:
             matches = re.findall(r"```sql(.*?)```", pred, re.DOTALL)
             return matches[-1].strip() if matches else pred
 
-    def make_predictions(self, prompts, model_name) -> list[dict]:
-        preds = []
-        for prompt in tqdm(prompts, desc=f"Analyzing Prompt with {model_name}"):
-            pred = self.make_prediction(prompt, model_name)
-            preds.append(pred)
-        return preds
 
-    def make_prediction(self, prompt, model_name):
+    def make_prediction(self, question, db_schema,  model_name, prompt=None):
         if model_name not in self.model_name2pred_func:
             raise ValueError(
                 "Model not supported",
@@ -82,32 +82,17 @@ class ModelPrediction:
                 self.model_name2pred_func.keys(),
             )
 
+        prompt = prompt or self.base_prompt
+        prompt = prompt.format(question=question, db_schema=db_schema)
+        print(prompt)
+
         prediction = self.model_name2pred_func[model_name](prompt)
         prediction["response_parsed"] = self._extract_answer_from_pred(
             prediction["response"]
         )
         return prediction
 
-    def _model_prediction(self, model_name):
-        predict_fun = self.predict_with_api
-        if "gpt-3.5" in model_name:
-            model_name = "openai/gpt-3.5-turbo-0125"
-        elif "gpt-4o-mini" in model_name:
-            model_name = "openai/gpt-4o-mini-2024-07-18"
-        elif "o1-mini" in model_name:
-            model_name = "openai/o1-mini-2024-09-12"
-        elif "QwQ" in model_name:
-            model_name = "together_ai/Qwen/QwQ-32B"
-        elif "DeepSeek-R1-Distill-Llama-70B" in model_name:
-            model_name = "together_ai/deepseek-ai/DeepSeek-R1-Distill-Llama-70B"
-        elif "llama-8" in model_name:
-            model_name = "meta-llama/Meta-Llama-3-8B-Instruct"
-            predict_fun = self.predict_with_hf
-        else:
-            raise ValueError("Model forbidden")
-
-        return partial(predict_fun, model_name=model_name)
-
+   
     def predict_with_api(self, prompt, model_name):  # -> dict[str, Any | float]:
         response = litellm.completion(
             model=model_name,
@@ -127,3 +112,24 @@ class ModelPrediction:
             "generated_text"
         ][-1]["content"]
         return {"response": response, "cost": 0.0}
+
+    def _init_model_prediction(self, model_name):
+        predict_fun = self.predict_with_api
+        if "gpt-3.5" in model_name:
+            model_name = "openai/gpt-3.5-turbo-0125"
+        elif "gpt-4o-mini" in model_name:
+            model_name = "openai/gpt-4o-mini-2024-07-18"
+        elif "o1-mini" in model_name:
+            model_name = "openai/o1-mini-2024-09-12"
+        elif "QwQ" in model_name:
+            model_name = "together_ai/Qwen/QwQ-32B"
+        elif "DeepSeek-R1-Distill-Llama-70B" in model_name:
+            model_name = "together_ai/deepseek-ai/DeepSeek-R1-Distill-Llama-70B"
+        elif "llama-8" in model_name:
+            model_name = "meta-llama/Meta-Llama-3-8B-Instruct"
+            predict_fun = self.predict_with_hf
+        else:
+            raise ValueError("Model forbidden")
+
+        return partial(predict_fun, model_name=model_name)
+

test_prediction.py

@@ -3,7 +3,7 @@ from prediction import ModelPrediction
 
 def main():
     model = ModelPrediction()
-    response = model.make_prediction("Hi, how are you?", "llama-8")
+    response = model.make_prediction(question='What is the name of Simone', db_schema='CREATE TABLE Player(Name, Age)', model_name="gpt-3.5", prompt='{question} {db_schema}')
     print(response)  # dict[response, response_parsed, cost]