Spaces:
Runtime error
Runtime error
Update app.py
Browse files
app.py
CHANGED
@@ -7,17 +7,26 @@ import spaces
|
|
7 |
# Load the model and tokenizer
|
8 |
peft_model_id = "rootxhacker/CodeAstra-7B"
|
9 |
config = PeftConfig.from_pretrained(peft_model_id)
|
10 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
11 |
tokenizer = AutoTokenizer.from_pretrained(config.base_model_name_or_path)
|
12 |
|
13 |
# Load the Lora model
|
14 |
model = PeftModel.from_pretrained(model, peft_model_id)
|
15 |
|
16 |
-
|
17 |
-
|
18 |
@spaces.GPU(duration=200)
|
19 |
def get_completion(query, model, tokenizer):
|
20 |
-
inputs = tokenizer(query, return_tensors="pt")
|
21 |
outputs = model.generate(**inputs, max_new_tokens=512, do_sample=True, temperature=0.7)
|
22 |
return tokenizer.decode(outputs[0], skip_special_tokens=True)
|
23 |
|
|
|
7 |
# Load the model and tokenizer
|
8 |
peft_model_id = "rootxhacker/CodeAstra-7B"
|
9 |
config = PeftConfig.from_pretrained(peft_model_id)
|
10 |
+
|
11 |
+
# Determine the device
|
12 |
+
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
13 |
+
|
14 |
+
# Load the model on the appropriate device
|
15 |
+
model = AutoModelForCausalLM.from_pretrained(
|
16 |
+
config.base_model_name_or_path,
|
17 |
+
return_dict=True,
|
18 |
+
load_in_4bit=True,
|
19 |
+
device_map="auto" # This will automatically handle device placement
|
20 |
+
)
|
21 |
+
|
22 |
tokenizer = AutoTokenizer.from_pretrained(config.base_model_name_or_path)
|
23 |
|
24 |
# Load the Lora model
|
25 |
model = PeftModel.from_pretrained(model, peft_model_id)
|
26 |
|
|
|
|
|
27 |
@spaces.GPU(duration=200)
|
28 |
def get_completion(query, model, tokenizer):
|
29 |
+
inputs = tokenizer(query, return_tensors="pt").to(device) # Move inputs to the same device as the model
|
30 |
outputs = model.generate(**inputs, max_new_tokens=512, do_sample=True, temperature=0.7)
|
31 |
return tokenizer.decode(outputs[0], skip_special_tokens=True)
|
32 |
|