implement tokenizer

main.py

@@ -5,9 +5,11 @@ from fastapi import FastAPI
 from pydantic import BaseModel
 from huggingface_hub import InferenceClient
 from typing import List
+import tiktoken
 
 app = FastAPI()
 client = InferenceClient("openai-community/gpt2-medium")
+tokenizer = tiktoken.get_encoding("gpt2")
 
 SYSTEM_PROMPT = (
     "You are a very powerful AI to generate interesting stories for short-form content "
@@ -15,6 +17,8 @@ SYSTEM_PROMPT = (
     "Make sure to be engaging and creative in your responses."
 )
 
+MAX_CONTEXT_LENGTH = 1024
+
 class Item(BaseModel):
     prompt: str
     history: List[str] = []
@@ -23,7 +27,6 @@ class Item(BaseModel):
     top_p: float = 0.15
     repetition_penalty: float = 1.0
 
-
 def format_prompt(message, history):
     prompt = "<s>"
     for user_prompt, bot_response in history:
@@ -31,16 +34,21 @@ def format_prompt(message, history):
     prompt += f"[INST] {message} [/INST]"
     return prompt
 
+def count_tokens(text):
+    return len(tokenizer.encode(text))
 
 def generate(item: Item):
     temperature = max(float(item.temperature), 1e-2)
 
     formatted_prompt = format_prompt(f"{SYSTEM_PROMPT}, {item.prompt}", item.history)
     
-    input_token_length = len(formatted_prompt.split())
-    max_allowed_tokens = 1024 - input_token_length 
+    input_token_length = count_tokens(formatted_prompt)
+    max_allowed_tokens = MAX_CONTEXT_LENGTH - input_token_length
     max_new_tokens = min(item.max_new_tokens, max_allowed_tokens)
 
+    if max_new_tokens <= 0:
+        raise ValueError("The input is too long. Please reduce the prompt or history length.")
+
     stream = client.text_generation(
         formatted_prompt,
         temperature=temperature,
@@ -60,13 +68,12 @@ def generate(item: Item):
 
     return output
 
-
 @app.get("/generate/")
 async def generate_text(
     prompt: str,
     history: List[str] = [],
     temperature: float = 0.0,
-    max_new_tokens: int = 1048,
+    max_new_tokens: int = 1024,
     top_p: float = 0.15,
     repetition_penalty: float = 1.0,
 ):
@@ -79,6 +86,11 @@ async def generate_text(
         repetition_penalty=repetition_penalty,
     )
 
-    response = await asyncio.to_thread(generate, item)
+    try:
+        response = await asyncio.to_thread(generate, item)
+        return {"response": response}
+    except ValueError as e:
+        return {"error": str(e)}
 
-    return {"response": response}
+if __name__ == "__main__":
+    uvicorn.run(app, host="0.0.0.0", port=8000)

requirements.txt

@@ -1,4 +1,5 @@
 fastapi
 uvicorn
 huggingface_hub
-pydantic
+pydantic
+tiktoken