Update app.py

app.py

@@ -10,34 +10,26 @@ model_id = "01-ai/Yi-34B-200K"
 os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'max_split_size_mb:54'
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-# Load the model and tokenizer using transformers
 tokenizer = YiTokenizer(vocab_file="./tokenizer.model")
-model = AutoModelForCausalLM.from_pretrained(model_id, trust_remote_code=True)
-model = model.to(torch.bfloat16)
+model = AutoModelForCausalLM.from_pretrained(model_id, load_in_8bit=True, trust_remote_code=True)
+# model = model.to(torch.bfloat16)
 model = model.to(device)
 
 def run(message, chat_history, max_new_tokens=4056, temperature=3.5, top_p=0.9, top_k=800):
     prompt = get_prompt(message, chat_history)
-
-    # Encode the prompt to tensor
     input_ids = tokenizer.encode(prompt, return_tensors='pt')
-
-    # Move input_ids to the same device as the model
     input_ids = input_ids.to(model.device)
-    
-    # Generate a response using the model with adjusted parameters
     response_ids = model.generate(
         input_ids,
         max_length=max_new_tokens + input_ids.shape[1],
-        temperature=temperature,  # Controls randomness. Lower values make text more deterministic.
-        top_p=top_p,              # Nucleus sampling: higher values allow more diversity.
-        top_k=top_k,              # Top-k sampling: limits the number of top tokens considered.
+        temperature=temperature,  
+        top_p=top_p,              
+        top_k=top_k,              
         pad_token_id=tokenizer.eos_token_id,
-        do_sample=True            # Enable sampling-based generation
+        do_sample=True            
 
     )
 
-    # Decode the response
     response = tokenizer.decode(response_ids[:, input_ids.shape[-1]:][0], skip_special_tokens=True)
     return response