Update agent.py
Browse files
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.
|
26 |
|
27 |
-
|
28 |
-
|
29 |
-
|
30 |
-
|
31 |
-
|
|
|
|
|
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")
|