Update app.py

app.py

@@ -8,20 +8,19 @@ access_token = os.getenv("read_access")
 from transformers import AutoModelForCausalLM, AutoTokenizer
 device = "cpu" # the device to load the model onto
 
-#tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2-0.5B-Instruct")
+tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2-0.5B-Instruct")
 
-#model1 = AutoModelForCausalLM.from_pretrained(
-#    "Qwen/Qwen2-1.5B-Instruct",
-#    device_map="auto"
-#)
-
-tokenizer = AutoTokenizer.from_pretrained("google/gemma-2-2b-it", token=access_token)
+model1 = AutoModelForCausalLM.from_pretrained(
+    "Qwen/Qwen2-1.5B-Instruct",
+    device_map="auto"
+)
 model = AutoModelForCausalLM.from_pretrained(
-    "google/gemma-2-2b-it",
-    device_map="auto",
-    token=access_token
+    "Qwen/Qwen2-1.5B-Instruct",
+    device_map="auto"
+    torch_dtype="auto"
 )
 
+
 app = FastAPI()
 
 @app.get("/")
@@ -77,7 +76,7 @@ async def read_droot():
         
     generated_ids = model1.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)
@@ -91,34 +90,3 @@ async def read_droot():
     return {"Hello": "World!"}
     #return {response: time}
 
-
-@app.get("/tet")
-async def read_droot():
-    starttime = time.time()
-    messages = [
-        {"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": "resps"}
-    #return {response: time}