Spaces:
Running
on
Zero
Running
on
Zero
Update app.py
Browse files
app.py
CHANGED
@@ -45,8 +45,8 @@ CHAT_TEMPLATE="""{%- set default_system_message = "A user will ask you to solve
|
|
45 |
|
46 |
if torch.cuda.is_available():
|
47 |
model_id = "mistralai/Mistral-Small-24B-Instruct-2501"
|
48 |
-
model = AutoModelForCausalLM.from_pretrained(
|
49 |
-
tokenizer = AutoTokenizer.from_pretrained(
|
50 |
|
51 |
|
52 |
@spaces.GPU
|
@@ -61,7 +61,7 @@ def generate(
|
|
61 |
) -> Iterator[str]:
|
62 |
conversation = [*chat_history, {"role": "user", "content": message}]
|
63 |
|
64 |
-
input_ids = tokenizer.apply_chat_template(conversation,
|
65 |
if input_ids.shape[1] > MAX_INPUT_TOKEN_LENGTH:
|
66 |
input_ids = input_ids[:, -MAX_INPUT_TOKEN_LENGTH:]
|
67 |
gr.Warning(f"Trimmed input from conversation as it was longer than {MAX_INPUT_TOKEN_LENGTH} tokens.")
|
|
|
45 |
|
46 |
if torch.cuda.is_available():
|
47 |
model_id = "mistralai/Mistral-Small-24B-Instruct-2501"
|
48 |
+
model = AutoModelForCausalLM.from_pretrained("AlexHung29629/fix_magistra4b", torch_dtype=torch.bfloat16, device_map="auto")
|
49 |
+
tokenizer = AutoTokenizer.from_pretrained("AlexHung29629/fix_magistra4")
|
50 |
|
51 |
|
52 |
@spaces.GPU
|
|
|
61 |
) -> Iterator[str]:
|
62 |
conversation = [*chat_history, {"role": "user", "content": message}]
|
63 |
|
64 |
+
input_ids = tokenizer.apply_chat_template(conversation, return_tensors="pt")
|
65 |
if input_ids.shape[1] > MAX_INPUT_TOKEN_LENGTH:
|
66 |
input_ids = input_ids[:, -MAX_INPUT_TOKEN_LENGTH:]
|
67 |
gr.Warning(f"Trimmed input from conversation as it was longer than {MAX_INPUT_TOKEN_LENGTH} tokens.")
|