Update app.py

app.py

@@ -4,10 +4,11 @@ import transformers
 from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig,AwqConfig
 import torch
 import os
+import bitnet
 key =  os.environ.get("key")
 from huggingface_hub import login 
 login(key)
-
+from bitnet import replace_linears_in_hf
 
 nf4_config = BitsAndBytesConfig(
    load_in_4bit=True,
@@ -16,15 +17,17 @@ nf4_config = BitsAndBytesConfig(
    bnb_4bit_compute_dtype=torch.bfloat16
 )
 
-model_id = "CohereForAI/c4ai-command-r-v01"
+model_id = "Nexusflow/Starling-LM-7B-beta"
 tokenizer = AutoTokenizer.from_pretrained(model_id)
 model = AutoModelForCausalLM.from_pretrained(model_id,
                                              # load_in_8bit=True,
-                                             quantization_config=nf4_config,
-                                             # torch_dtype = torch.bfloat16,
+                                             # quantization_config=nf4_config,
+                                             torch_dtype = torch.bfloat16,
                                              # device_map="auto"
                                             )
 
+replace_linears_in_hf(model)
+model..to('cuda').eval()
 @spaces.GPU
 def generate_response(user_input, max_new_tokens, temperature):
     os.system("nvidia-smi")