Coool2 commited on
Commit
f0ebeea
·
1 Parent(s): 8b1644b

Update agent.py

Browse files
Files changed (1) hide show
  1. agent.py +21 -22
agent.py CHANGED
@@ -120,28 +120,27 @@ def initialize_models(use_api_mode=False):
120
  # Try to use Pixtral 12B with vLLM if available
121
  pixtral_model = "Qwen/Qwen3-8B" # Fallback model
122
  try:
123
- if os.path.exists("/path/to/pixtral-12b") or True: # Placeholder check
124
- pixtral_model = "mistralai/pixtral-12b"
125
- print(f"Using Pixtral 12B with vLLM")
126
-
127
- # Custom prompt template for Pixtral model
128
- def messages_to_prompt(messages):
129
- prompt = "\n".join([str(x) for x in messages])
130
- return f"<s>[INST] {prompt} [/INST] </s>\n"
131
-
132
- def completion_to_prompt(completion):
133
- return f"<s>[INST] {completion} [/INST] </s>\n"
134
-
135
- proj_llm = Vllm(
136
- model=pixtral_model,
137
- tensor_parallel_size=1, # Adjust based on available GPUs
138
- max_new_tokens=16000,
139
- messages_to_prompt=messages_to_prompt,
140
- completion_to_prompt=completion_to_prompt,
141
- temperature=0.6,
142
- top_p=0.95,
143
- top_k=20
144
- )
145
  else:
146
  # Use regular Qwen model if Pixtral not found
147
  raise ImportError("Pixtral 12B not found")
 
120
  # Try to use Pixtral 12B with vLLM if available
121
  pixtral_model = "Qwen/Qwen3-8B" # Fallback model
122
  try:
123
+ pixtral_model = "mistralai/Pixtral-12B-2409"
124
+ print(f"Using Pixtral 12B with vLLM")
125
+
126
+ # Custom prompt template for Pixtral model
127
+ def messages_to_prompt(messages):
128
+ prompt = "\n".join([str(x) for x in messages])
129
+ return f"<s>[INST] {prompt} [/INST] </s>\n"
130
+
131
+ def completion_to_prompt(completion):
132
+ return f"<s>[INST] {completion} [/INST] </s>\n"
133
+
134
+ proj_llm = Vllm(
135
+ model=pixtral_model,
136
+ tensor_parallel_size=1, # Adjust based on available GPUs
137
+ max_new_tokens=16000,
138
+ messages_to_prompt=messages_to_prompt,
139
+ completion_to_prompt=completion_to_prompt,
140
+ temperature=0.6,
141
+ top_p=0.95,
142
+ top_k=20
143
+ )
 
144
  else:
145
  # Use regular Qwen model if Pixtral not found
146
  raise ImportError("Pixtral 12B not found")