Update main.py

main.py

@@ -21,12 +21,12 @@ model_id = "google/gemma-2-2b-it"
 device = "cuda:0" if torch.cuda.is_available() else "cpu"
 
 # Load tokenizer and model with authentication token
-tokenizer = AutoTokenizer.from_pretrained(model_id, use_auth_token=HF_TOKEN)
+tokenizer = AutoTokenizer.from_pretrained(model_id, token=HF_TOKEN)
 model = AutoModelForCausalLM.from_pretrained(
     model_id,
     device_map="auto",
     torch_dtype=torch.float16,
-    use_auth_token=HF_TOKEN
+    token=HF_TOKEN
 )
 
 app_pipeline = pipeline(
@@ -44,6 +44,7 @@ def generate_text():
     temperature = data.get("temperature", 0.1)
     top_k = data.get("top_k", 50)
     top_p = data.get("top_p", 0.95)
+    print(f"{prompt}: ")
 
     try:
         outputs = app_pipeline(