Update app.py

app.py

@@ -7,45 +7,31 @@ 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 = FastAPI()
 
 @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."},
@@ -53,32 +39,20 @@ messages = [
         {"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 
@@ -103,7 +77,7 @@ async def read_droot():
         
     generated_ids = model.generate(
         model_inputs.input_ids,
-        max_new_tokens=64
+        max_new_tokens=128
     )
     generated_ids = [
         output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
@@ -116,37 +90,6 @@ async def read_droot():
     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()
@@ -161,7 +104,7 @@ async def read_droot():
         tokenize=False,
         add_generation_prompt=True
     )
-    model_inputs = tokenizer1([text], return_tensors="pt").to(device)
+    model_inputs = tokenizer([text], return_tensors="pt").to(device)
         
     generated_ids = model1.generate(
         model_inputs.input_ids,
@@ -171,7 +114,7 @@ async def read_droot():
         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]
+    response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
     print(response)
     end_time = time.time()
     time_taken = end_time - starttime 
@@ -205,7 +148,7 @@ async def read_droot():
     ]
         
     response = tokenizer2.batch_decode(generated_ids, skip_special_tokens=True)[0]
-    respons = tokenizer1.batch_decode(generated_ids, skip_special_tokens=True)[0]
+    respons = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
     print(response)
     end_time = time.time()
     time_taken = end_time - starttime