reab5555 commited on
Commit
c5a0a0b
·
verified ·
1 Parent(s): 9763663

Update llm_loader.py

Browse files
Files changed (1) hide show
  1. llm_loader.py +9 -2
llm_loader.py CHANGED
@@ -1,10 +1,17 @@
1
  from langchain_openai import ChatOpenAI
 
 
 
2
 
3
  def load_model(openai_api_key):
4
  return ChatOpenAI(
5
- model_name="gpt-4o-mini",
6
  openai_api_key=openai_api_key,
7
  temperature=0.01,
8
  max_tokens=4096,
9
  top_p=0.9
10
- )
 
 
 
 
 
1
  from langchain_openai import ChatOpenAI
2
+ from tiktoken import encoding_for_model
3
+
4
+ model = "gpt-4o-mini"
5
 
6
  def load_model(openai_api_key):
7
  return ChatOpenAI(
8
+ model_name=model,
9
  openai_api_key=openai_api_key,
10
  temperature=0.01,
11
  max_tokens=4096,
12
  top_p=0.9
13
+ )
14
+
15
+ def count_tokens(text, model=model):
16
+ encoder = encoding_for_model(model)
17
+ return len(encoder.encode(text))