Update app.py

app.py

@@ -2,6 +2,8 @@ import os
 import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM
 from huggingface_hub import login
+import time
+import torch.quantization
 
 # Directly assign your Hugging Face token here
 hf_token = "your_hugging_face_api_token"
@@ -19,24 +21,49 @@ try:
 except ImportError:
     raise ImportError("The sentencepiece library is required for this tokenizer. Please install it with `pip install sentencepiece`.")
 
+# Start time to measure execution time
+start_time = time.time()
+
+# Load tokenizer and model
 tokenizer = AutoTokenizer.from_pretrained(model_name, use_auth_token=hf_token)
 model = AutoModelForCausalLM.from_pretrained(model_name, use_auth_token=hf_token)
 
+# Quantize the model
+quantized_model = torch.quantization.quantize_dynamic(model, {torch.nn.Linear}, dtype=torch.qint8)
+
 # Check if a GPU is available and if not, fall back to CPU
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-model.to(device)
+quantized_model.to(device)
+
+# Measure time for loading tokenizer, model, and quantization
+loading_time = time.time() - start_time
+print(f"Time taken to load tokenizer, model, and quantize: {loading_time:.2f} seconds")
 
 # Example text input
 text_input = "How did Tesla perform in Q1 2024?"
 
+# Start time for inference
+inference_start_time = time.time()
+
 # Tokenize the input text
 inputs = tokenizer(text_input, return_tensors="pt").to(device)
 
-# Generate a response without sampling (deterministic)
-outputs = model.generate(**inputs, max_length=150, do_sample=False)
+# Measure time for tokenization
+tokenization_time = time.time() - inference_start_time
+
+# Generate a response
+outputs = quantized_model.generate(**inputs, max_length=150, do_sample=False)
+
+# Measure time for inference
+inference_time = time.time() - inference_start_time
+print(f"Time taken for inference with quantized model: {inference_time:.2f} seconds")
 
 # Decode the generated tokens to a readable string
 response = tokenizer.decode(outputs[0], skip_special_tokens=True)
 
 # Print the response
-print(response)
+print(f"Generated response: {response}")
+
+# Total execution time
+total_time = time.time() - start_time
+print(f"Total execution time with quantized model: {total_time:.2f} seconds")