roll back to gpt-2

main.py

@@ -5,54 +5,49 @@ 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")
+client = InferenceClient("openai-community/gpt2")
 
-SYSTEM_PROMPT = (
-    "You are a very powerful AI to generate interesting stories for short-form content "
-    "consumption. Make sure to hook the reader's attention in the first few seconds. "
-    "Make sure to be engaging and creative in your responses."
-)
+SYSTEM_PROMPT = "You are a very powerful AI to generate interesting stories for short-form content consumption. Make sure to hook the readers attention in the first few seconds. Make sure to be engaging and creative in your responses."
 
-MAX_CONTEXT_LENGTH = 1024
 
 class Item(BaseModel):
     prompt: str
     history: List[str] = []
+    # system_prompt: str = "You are a very powerful AI assistant."
     temperature: float = 0.0
     max_new_tokens: int = 1024
     top_p: float = 0.15
     repetition_penalty: float = 1.0
 
+
 def format_prompt(message, history):
     prompt = "<s>"
     for user_prompt, bot_response in history:
-        prompt += f"[INST] {user_prompt} [/INST] {bot_response}</s> "
+        prompt += f"[INST] {user_prompt} [/INST]"
+        prompt += f" {bot_response}</s> "
     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)
+    # generate_kwargs = dict(
+    #     temperature=temperature,
+    #     max_new_tokens=item.max_new_tokens,
+    #     top_p=float(item.top_p),
+    #     repetition_penalty=item.repetition_penalty,
+    #     do_sample=True,
+    #     seed=42,
+    # )
 
     formatted_prompt = format_prompt(f"{SYSTEM_PROMPT}, {item.prompt}", item.history)
-    
-    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,
-        max_new_tokens=max_new_tokens,
+        max_new_tokens=item.max_new_tokens,
         top_p=float(item.top_p),
         repetition_penalty=item.repetition_penalty,
         do_sample=True,
@@ -61,36 +56,34 @@ def generate(item: Item):
         details=True,
         return_full_text=False,
     )
-
     output = "".join(response.token.text for response in stream)
-    output = re.sub(r"<[^>]+>", "", output)
-    output = re.sub(r"\s+", " ", output).strip()
+    # Remove unwanted sequences or patterns (e.g., <s>, [/INST], etc.)
+    output = re.sub(r"<[^>]+>", "", output)  # Remove any HTML-like tags
+    output = re.sub(r"\s+", " ", output).strip()  # Clean up extra whitespace
 
     return output
 
+
 @app.get("/generate/")
 async def generate_text(
     prompt: str,
     history: List[str] = [],
+    # system_prompt: str = "You are a very powerful AI assistant.",
     temperature: float = 0.0,
-    max_new_tokens: int = 1024,
+    max_new_tokens: int = 1048,
     top_p: float = 0.15,
     repetition_penalty: float = 1.0,
 ):
     item = Item(
         prompt=prompt,
         history=history,
+        # system_prompt=system_prompt,
         temperature=temperature,
         max_new_tokens=max_new_tokens,
         top_p=top_p,
         repetition_penalty=repetition_penalty,
     )
 
-    try:
-        response = await asyncio.to_thread(generate, item)
-        return {"response": response}
-    except ValueError as e:
-        return {"error": str(e)}
+    response = await asyncio.to_thread(generate, item)
 
-if __name__ == "__main__":
-    uvicorn.run(app, host="0.0.0.0", port=8000)
+    return {"response": response}