Update main.py

main.py

@@ -1,14 +1,15 @@
 from fastapi import FastAPI
 from pydantic import BaseModel
-from huggingface_hub import InferenceClient
 import uvicorn
+import prompt_style
+
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import torch
 
 # **************************************************
 # import transformers
 # import torch
 
-model_id = "meta-llama/Meta-Llama-3-8B-Instruct"
-
 # pipeline = transformers.pipeline(
 #     "text-generation",
 #     model=model_id,
@@ -16,7 +17,7 @@ model_id = "meta-llama/Meta-Llama-3-8B-Instruct"
 #     device_map="auto",
 # )
 
-def generate(item: Item):
+def generate_1(item: Item):
     messages = [
         {"role": "system", "content": "You are a pirate chatbot who always responds in pirate speak!"},
         {"role": "user", "content": "Who are you?"},
@@ -55,13 +56,14 @@ class Item(BaseModel):
     
 app = FastAPI()
 
-def format_prompt(message, history):
-    prompt = "<s>"
-    for user_prompt, bot_response in history:
-        prompt += f"[INST] {user_prompt} [/INST]"
-        prompt += f" {bot_response}</s> "
-    prompt += f"[INST] {message} [/INST]"
-    return prompt
+def format_prompt(item: Item):
+    messages = [
+        {"role": "system", "content": prompt_style.data},
+    ]
+    for it in item.history:
+        messages.append[{"role" : "user", "content": it[0]}]
+        messages.append[{"role" : "assistant", "content": it[1]}]
+    return messages
 
 def generate(item: Item):
     temperature = float(item.temperature)
@@ -78,13 +80,29 @@ def generate(item: Item):
         seed=item.seed,
     )
 
-    formatted_prompt = format_prompt(f"{item.system_prompt}, {item.prompt}", item.history)
-    stream = client.text_generation(formatted_prompt, **generate_kwargs, stream=True, details=True, return_full_text=False)
-    output = ""
+    formatted_prompt = format_prompt(item)
+
+    model_id = "meta-llama/Meta-Llama-3-8B-Instruct"
+    tokenizer = AutoTokenizer.from_pretrained(model_id)
+    model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.bfloat16, device_map="auto",)
+
+    input_ids = tokenizer.apply_chat_template(formatted_prompt, add_generation_prompt=True, return_tensors="pt").to(model.device)
+    
+    terminators = [
+        tokenizer.eos_token_id,
+        tokenizer.convert_tokens_to_ids("<|eot_id|>")
+    ]
+
+    outputs = model.generate(input_ids, eos_token_id=terminators, do_sample=True, **generate_kwargs)
+    response = outputs[0][input_ids.shape[-1]:]
+    return tokenizer.decode(response, skip_special_tokens=True)
+    
+    # stream = client.text_generation(formatted_prompt, **generate_kwargs, stream=True, details=True, return_full_text=False)
+    # output = ""
 
-    for response in stream:
-        output += response.token.text
-    return output
+    # for response in stream:
+    #     output += response.token.text
+    # return output
 
 @app.post("/generate/")
 async def generate_text(item: Item):