Added trace logs

main.py

@@ -12,6 +12,7 @@
 from fastapi import FastAPI
 from pydantic import BaseModel
 from transformers import AutoModelForCausalLM, AutoTokenizer
+import sys
 import uvicorn
 import torch
 # torch.mps.empty_cache()
@@ -67,20 +68,32 @@ def generate(item: Item):
     device = "cuda" if torch.cuda.is_available() else "cpu"
 
     model_path = "ibm-granite/granite-34b-code-instruct-8k"
+    
+    print("Loading tokenizer for model: " + model_path, file=sys.stderr)
     tokenizer = AutoTokenizer.from_pretrained(model_path, cache_dir="/code/huggingface/transformers")
     # drop device_map if running on CPU
+    
+    print("Loading Model for causal LM for model: " + model_path, file=sys.stderr)
     model = AutoModelForCausalLM.from_pretrained(model_path, device_map=device)
     model.eval()
     # change input text as desired
     chat = format_prompt(item.system_prompt, item.prompt, item.history)
     chat = tokenizer.apply_chat_template(chat, tokenize=False, add_generation_prompt=True)
+    
+    print("Tokenizing text", file=sys.stderr)
     # tokenize the text
     input_tokens = tokenizer(chat, return_tensors="pt")
+    
+    print("Transferring tokens to device: " + device, file=sys.stderr)
     # transfer tokenized inputs to the device
     for i in input_tokens:
         input_tokens[i] = input_tokens[i].to(device)
+        
+    print("Generating output tokens", file=sys.stderr)
     # generate output tokens
     output = model.generate(**input_tokens, max_new_tokens=900)
+    
+    print("Decoding output tokens", file=sys.stderr)
     output_text = tokenizer.batch_decode(output, skip_special_tokens=True)[0]
     return output_text