casperhansen

/

llama-3-70b-instruct-awq

like 66

Good day!
I am testing this model with transformers' pipeline, but an API response is too slow (~100s), could someone please take a look at my code, what am I doing wrong?
I am using NVIDIA A40 GPU server on RunPod (48GB VRAM, 48GB RAM).

import transformers
import torch
from fastapi import FastAPI
from pydantic import BaseModel

app = FastAPI()

pipeline = None

def load_model():
    global pipeline

    model_id = "casperhansen/llama-3-70b-instruct-awq"
    device_map = "auto"

    pipeline = transformers.pipeline(
    "text-generation",
    model=model_id,
    model_kwargs={"torch_dtype": torch.float16},
    device_map=device_map)

load_model()

class TextGenerationResponse(BaseModel):
    generated_text: str

class TextGenerationRequest(BaseModel):
    systemPrompt: str
    prompt: str
    temperature: float

def generate_inner(request: TextGenerationRequest):
    messages = [
        {"role": "system", "content": request.systemPrompt},
        {"role": "user", "content": request.prompt},
    ]
    
    prompt = pipeline.tokenizer.apply_chat_template(
        messages, 
        tokenize=False, 
        add_generation_prompt=True)

    terminators = [
    pipeline.tokenizer.eos_token_id,
    pipeline.tokenizer.convert_tokens_to_ids("<|eot_id|>")
    ]
    
    outputs = pipeline(
        prompt,
        max_new_tokens=512,
        eos_token_id=terminators,
        do_sample=True,
        temperature=request.temperature,
        top_p=0.9,
    )
    text = outputs[0]["generated_text"][len(prompt):]

    return text



@app

	.post("/generate")
def generate(request: TextGenerationRequest):

    if (request.temperature < 0.1 and request.temperature > 2): request.temperature = 0.6

    generatedText = generate_inner(request)
    
    return [
        TextGenerationResponse(generated_text=generatedText)
    ]