Spaces:
Configuration error
Configuration error
Update agent.py
Browse files
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 |
-
|
124 |
-
|
125 |
-
|
126 |
-
|
127 |
-
|
128 |
-
|
129 |
-
|
130 |
-
|
131 |
-
|
132 |
-
|
133 |
-
|
134 |
-
|
135 |
-
|
136 |
-
|
137 |
-
|
138 |
-
|
139 |
-
|
140 |
-
|
141 |
-
|
142 |
-
|
143 |
-
|
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")
|