feat add prediction with togheter AI and HF pipe

prediction.py

@@ -1,45 +1,127 @@
+from functools import partial
+import os
+import re
+from xml.parsers.expat import model
+
 # https://discuss.huggingface.co/t/issues-with-sadtalker-zerogpu-spaces-inquiry-about-community-grant/110625/10
 if os.environ.get("SPACES_ZERO_GPU") is not None:
     import spaces
 else:
+
     class spaces:
         @staticmethod
         def GPU(func):
             def wrapper(*args, **kwargs):
                 return func(*args, **kwargs)
+
             return wrapper
 
 
+from transformers import pipeline as hf_pipeline
+import torch
+import litellm
+
 
 class ModelPrediction:
     def __init__(self, model_name):
-        self.prediction_fun = self._model_prediction(model_name)
+        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(
+                "DeepSeek-R1-Distill-Llama-70B"
+            ),
+        }
+
+        self._model_name = None
+        self._pipeline = None
+
+    @property
+    def pipeline(self):
+        if self._pipeline is None:
+            self._pipeline = hf_pipeline(
+                task="text-generation",
+                model=self._model_name,
+                torch_dtype=torch.bfloat16,
+                device_map="auto",
+            )
+        return self._pipeline
+
+    def _reset_pipeline(self, model_name):
+        if self._model_name != model_name:
+            self._model_name = model_name
+            self._pipeline = None
+
+    @staticmethod
+    def _extract_answer_from_pred(pred: str) -> str:
+        # extract with regex everything is between <answer> and </answer>
+        matches = re.findall(r"<answer>(.*?)</answer>", pred, re.DOTALL)
+        if matches:
+            return matches[-1].replace("```", "").replace("sql", "").strip()
+        else:
+            matches = re.findall(r"```sql(.*?)```", pred, re.DOTALL)
+            return matches[-1].strip() if matches else pred
+
+    def make_prediction(self, prompt, model_name):
+        if model_name not in self.model_name2pred_func:
+            raise ValueError(
+                "Model not supported",
+                "supported models are",
+                self.model_name2pred_func.keys(),
+            )
+
+        prediction = self.model_name2pred_func[model_name](prompt)
+        prediction["response_parsed"] = self._extract_answer_from_pred(
+            prediction["response"]
+        )
+        return prediction
 
-    def make_prediction(prompt):
-        pass
-     
     def _model_prediction(self, model_name):
-        predict_fun = 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'
+        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"
         else:
-            raise ValueError('Model forbidden')
-        
-        return 
+            raise ValueError("Model forbidden")
 
-    def predict_with_api(prompt):
-        pass
+        return partial(predict_fun, model_name=model_name)
 
+    def predict_with_api(self, prompt, model_name):  # -> dict[str, Any | float]:
+        def track_cost_callback(
+            kwargs,  # kwargs to completion
+            completion_response,  # response from completion
+            start_time,
+            end_time,  # start/end time
+        ):
+            try:
+                response_cost = kwargs[
+                    "response_cost"
+                ]  # litellm calculates response cost for you
+                call_cost = response_cost
+            except:
+                pass
 
-    @spaces.GPU
-    def predict_with_hf(prompt):
-        pass
+        litellm.success_callback = [track_cost_callback]
+        call_cost = 0.0
+        response = litellm.completion(
+            model=model_name,
+            messages=[{"role": "user", "content": prompt}],
+            num_retries=2,
+        )
+        return {"response": response, "cost": call_cost}
 
+    @spaces.GPU
+    def predict_with_hf(self, prompt, model_name):  # -> dict[str, Any | float]:
+        self._reset_pipeline(model_name)
+        response = self.pipeline([{"role": "user", "content": prompt}])[0][
+            "generated_text"
+        ][-1]["content"]
+        return {"response": response, "cost": 0.0}

requirements.txt

@@ -10,7 +10,7 @@ eval-type-backport>=0.2.0
 openai==1.66.3
 litellm==1.63.14
 together==1.4.6
-
+litellm==1.63.14
 # Conditional dependency for Gradio (requires Python >=3.10)
 gradio>=5.20.1; python_version >= "3.10"