Attempting with hf model

app.py

@@ -8,11 +8,13 @@ from sse_starlette.sse import EventSourceResponse
 
 from config.log_config import uvicorn_logger
 from models import OpenAIinput
-from utils.codegen import CodeGenProxy
+from utils.hf_model import HFModel
 from utils.errors import FauxPilotException
 
 logging.config.dictConfig(uvicorn_logger)
 
+model = HFModel("bigcode/santacoder")
+
 codegen = CodeGenProxy(
     host=os.environ.get("TRITON_HOST", "triton"),
     port=os.environ.get("TRITON_PORT", 8001),
@@ -63,3 +65,11 @@ async def completions(data: OpenAIinput):
 
 if __name__ == "__main__":
     uvicorn.run("app:app", host="0.0.0.0", port=5000)
+
+
+# curl request to test the API
+curl -X POST "http://localhost:5000/v1/engines/codegen/completions" -H "accept: application/json" -H "Content-Type: application/json" -d "{\"prompt\": \"import numpy as np\"}"
+
+curl -X POST "https://huggingface.co/spaces/ncoop57/santacoder-openai/v1/engines/codegen/completions" -H "accept: application/json" -H "Content-Type: application/json" -d "{\"prompt\": \"import numpy as np\"}"
+
+curl -X POST "https://ncoop57-santacoder-openai.hf.space/v1/engines/codegen/completions" -H "accept: application/json" -H "Content-Type: application/json" -d "{\"prompt\": \"import numpy as np\"}"

utils/hf_model.py

@@ -0,0 +1,27 @@
+import time
+
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+class HFModel:
+    def __init__(self, model_name):
+        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
+        self.model = AutoModelForCausalLM.from_pretrained(model_name)
+
+    def __call__(self, request):
+        prompt = request.get("prompt")
+        input_ids = self.tokenizer.encode(prompt, return_tensors='pt')
+        choices = self.model.generate(input_ids, max_length=50, do_sample=True)
+        choices self.tokenizer.decode(choices, skip_special_tokens=True)
+        completion = {
+            'id': None,  # fill in
+            'model': 'codegen',
+            'object': 'text_completion',
+            'created': int(time.time()),
+            'choices': None,  # fill in
+            'usage': {
+                'completion_tokens': int(sum([len(c.split()) for c in choices])),
+                'prompt_tokens': int(len(prompt.split())),
+                'total_tokens': int(sum([len(c.split()) for c in choices]) + len(prompt.split())),
+            }
+        }
+        return completion, choices