Coool2 commited on
Commit
87d5798
·
verified ·
1 Parent(s): 410bf79

Update agent.py

Browse files
Files changed (1) hide show
  1. agent.py +8 -6
agent.py CHANGED
@@ -22,13 +22,15 @@ from llama_index.core.callbacks.llama_debug import LlamaDebugHandler
22
  from llama_index.core import Settings
23
 
24
  from transformers import AutoModelForCausalLM, AutoTokenizer
25
- from llama_index.llms.vllm import Vllm
26
 
27
- llm = Vllm(
28
- model="mistralai/Pixtral-12B-2409",
29
- tensor_parallel_size=2, # For two GPUs
30
- max_new_tokens=512,
31
- vllm_kwargs={"swap_space": 1, "gpu_memory_utilization": 0.9},
 
 
32
  )
33
 
34
  embed_model = HuggingFaceEmbedding("BAAI/bge-small-en-v1.5")
 
22
  from llama_index.core import Settings
23
 
24
  from transformers import AutoModelForCausalLM, AutoTokenizer
25
+ from llama_index.llms.huggingface import HuggingFaceLLM
26
 
27
+ model_id = "Qwen/Qwen2.5-14B-Instruct"
28
+ proj_llm = HuggingFaceLLM(
29
+ model_name=model_id,
30
+ tokenizer_name=model_id,
31
+ device_map="auto", # will use GPU if available
32
+ model_kwargs={"torch_dtype": "auto"},
33
+ generate_kwargs={"temperature": 0.7, "top_p": 0.95}
34
  )
35
 
36
  embed_model = HuggingFaceEmbedding("BAAI/bge-small-en-v1.5")