Update app.py

app.py

@@ -3,7 +3,7 @@ from peft import PeftModel, PeftConfig
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
 
-
+'''
 if torch.cuda.is_available():
     device = torch.device("cuda")
     print("GPU is available!")
@@ -19,13 +19,38 @@ tokenizer = AutoTokenizer.from_pretrained(config.base_model_name_or_path)
 
 # Load the Lora model
 model = PeftModel.from_pretrained(model, peft_model_id)
+'''
+
+import torch
+from peft import PeftModel, PeftConfig
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+if torch.cuda.is_available():
+    device = torch.device("cuda")
+    print("GPU is available!")
+else:
+    device = torch.device("cpu")
+    print("GPU is not available, using CPU.")
+
+peft_model_id = "phearion/bigbrain-v0.0.1"
+config = PeftConfig.from_pretrained(peft_model_id)
+model = AutoModelForCausalLM.from_pretrained(
+    config.base_model_name_or_path,                          
+    low_cpu_mem_usage=True,
+    return_dict=True,
+    torch_dtype=torch.bfloat16)
+
+tokenizer = AutoTokenizer.from_pretrained(config.base_model_name_or_path)
+
+# Load the Lora model
+model = PeftModel.from_pretrained(model, peft_model_id)
+model = model.merge_and_unload()
 
 def greet(text):
-    batch = tokenizer(f"\"{text}\" ->: ", return_tensors='pt')
-    
-    # Use torch.no_grad to disable gradient calculation
-    with torch.no_grad():
-        output_tokens = model.generate(**batch, do_sample=True, max_new_tokens=20)
+    batch = tokenizer("“aide moi avec les equa diff ” ->: ", return_tensors='pt')
+
+    with torch.cuda.amp.autocast():
+      output_tokens = model.generate(**batch, max_new_tokens=15)
 
     return tokenizer.decode(output_tokens[0], skip_special_tokens=True)