Update app.py

app.py

@@ -3,7 +3,7 @@ import torch
 import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM, AutoModelForSequenceClassification
 
-model_id = "nicholasKluge/Aira-Instruct-124M" 
+model_id = "nicholasKluge/Aira-Instruct-1.5B" # "nicholasKluge/Aira-Instruct-124M"
 rewardmodel_id = "nicholasKluge/RewardModel"
 toxicitymodel_id = "nicholasKluge/ToxicityModel"
 
@@ -11,7 +11,9 @@ token = "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"
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-model = AutoModelForCausalLM.from_pretrained(model_id, use_auth_token=token)
+model = AutoModelForCausalLM.from_pretrained(model_id, use_auth_token=token, device_map="auto", load_in_8bit=True)
+
+#model = AutoModelForCausalLM.from_pretrained(model_id, use_auth_token=token)
 rewardModel = AutoModelForSequenceClassification.from_pretrained(rewardmodel_id, use_auth_token=token)
 toxicityModel = AutoModelForSequenceClassification.from_pretrained(toxicitymodel_id, use_auth_token=token)
 
@@ -19,7 +21,7 @@ model.eval()
 rewardModel.eval()
 toxicityModel.eval()
 
-model.to(device)
+#model.to(device)
 rewardModel.to(device)
 toxicityModel.to(device)