app.py

from fastapi import FastAPI
import time
import torch
import os

access_token = os.getenv("read_access")

from transformers import AutoModelForCausalLM, AutoTokenizer
device = "cpu" # the device to load the model onto
time1 = time.time()
model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen2-0.5B-Instruct",
    torch_dtype="auto",
    device_map="auto"
)
time2 = time.time()
print(time2-time1)
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2-0.5B-Instruct")
time3 = time.time()
print(time3-time1)
model1 = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen2-1.5B-Instruct",
    torch_dtype="auto",
    device_map="auto"
)
tokenizer1 = AutoTokenizer.from_pretrained("Qwen/Qwen2-1.5B-Instruct")
time4 = time.time()
print(time4-time3)
app = FastAPI()
time5 = time.time()
print(time5-time4)


tokenizer2 = AutoTokenizer.from_pretrained("google/gemma-2-2b-it", token=access_token)
model2 = AutoModelForCausalLM.from_pretrained(
    "google/gemma-2-2b-it",
    device_map="auto",
    torch_dtype=torch.bfloat16,
    token=access_token
)
model3 = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen2-0.5B-Instruct",
    device_map="auto"
)

@app.get("/")
async def read_root():
    return {"Hello": "World!"}
start_time = time.time()
messages = [
        {"role": "system", "content": "You are a helpful assistant, Sia, developed by Sushma. You will response in polity and brief."},
        {"role": "user", "content": "I'm Alok. Who are you?"},
        {"role": "assistant", "content": "I am Sia, a small language model created by Sushma."},
        {"role": "user", "content": "How are you?"}
]
time1 = time.time()
text = tokenizer.apply_chat_template(
    messages,
    tokenize=False,
    add_generation_prompt=True
)
time2 = time.time()
print(time2-time1)
model_inputs = tokenizer([text], return_tensors="pt").to(device)
time3 = time.time()
print(time3-time2)    
generated_ids = model.generate(
    model_inputs.input_ids,
    max_new_tokens=64
)
time4 = time.time()
print(time4-time3)
generated_ids = [
    output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
]
time5 = time.time()
print(time5-time4)    
response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
print(response)
time6 = time.time()
print(time6-time5)

end_time = time.time()
time_taken = end_time - start_time 
print(time_taken)


@app.get("/test")
async def read_droot():
    starttime = time.time()
    messages = [
        {"role": "system", "content": "You are a helpful assistant, Sia, developed by Sushma. You will response in polity and brief."},
        {"role": "user", "content": "I'm Alok. Who are you?"},
        {"role": "assistant", "content": "I am Sia, a small language model created by Sushma."},
        {"role": "user", "content": "How are you?"}
    ]
    text = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True
    )
    model_inputs = tokenizer([text], return_tensors="pt").to(device)
        
    generated_ids = model.generate(
        model_inputs.input_ids,
        max_new_tokens=64
    )
    generated_ids = [
        output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
    ]
        
    response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
    print(response)
    end_time = time.time()
    time_taken = end_time - starttime 
    print(time_taken)
    return {"Hello": "World!"}

@app.get("/teat")
async def read_droot():
    starttime = time.time()
    messages = [
        {"role": "system", "content": "You are a helpful assistant, Sia, developed by Sushma. You will response in polity and brief."},
        {"role": "user", "content": "I'm Alok. Who are you?"},
        {"role": "assistant", "content": "I am Sia, a small language model created by Sushma."},
        {"role": "user", "content": "How are you?"}
    ]
    text = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True
    )
    model_inputs = tokenizer([text], return_tensors="pt").to(device)
        
    generated_ids = model3.generate(
        model_inputs.input_ids,
        max_new_tokens=64
    )
    generated_ids = [
        output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
    ]
        
    response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
    print(response)
    end_time = time.time()
    time_taken = end_time - starttime 
    print(time_taken)
    return {"Hello": "World!"}    

@app.get("/text")
async def read_droot():
    starttime = time.time()
    messages = [
        {"role": "system", "content": "You are a helpful assistant, Sia, developed by Sushma. You will response in polity and brief."},
        {"role": "user", "content": "I'm Alok. Who are you?"},
        {"role": "assistant", "content": "I am Sia, a small language model created by Sushma."},
        {"role": "user", "content": "How are you?"}
    ]
    text = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True
    )
    model_inputs = tokenizer1([text], return_tensors="pt").to(device)
        
    generated_ids = model1.generate(
        model_inputs.input_ids,
        max_new_tokens=64
    )
    generated_ids = [
        output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
    ]
        
    response = tokenizer1.batch_decode(generated_ids, skip_special_tokens=True)[0]
    print(response)
    end_time = time.time()
    time_taken = end_time - starttime 
    print(time_taken)
    return {"Hello": "World!"}
    #return {response: time}


@app.get("/tet")
async def read_droot():
    starttime = time.time()
    messages = [
        {"role": "system", "content": "You are a helpful assistant, Sia, developed by Sushma. You will response in polity and brief."},
        {"role": "user", "content": "I'm Alok. Who are you?"},
        {"role": "assistant", "content": "I am Sia, a small language model created by Sushma."},
        {"role": "user", "content": "How are you?"}
    ]
    text = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True
    )
    model_inputs = tokenizer2([text], return_tensors="pt").to(device)
        
    generated_ids = model2.generate(
        model_inputs.input_ids,
        max_new_tokens=64
    )
    generated_ids = [
        output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
    ]
        
    response = tokenizer2.batch_decode(generated_ids, skip_special_tokens=True)[0]
    respons = tokenizer1.batch_decode(generated_ids, skip_special_tokens=True)[0]
    print(response)
    end_time = time.time()
    time_taken = end_time - starttime 
    print(time_taken)
    return {"Hello": respons}
    #return {response: time}