VyLala commited on
Commit
d3e0e88
·
verified ·
1 Parent(s): 42719fb

Update model.py

Browse files
Files changed (1) hide show
  1. model.py +19 -11
model.py CHANGED
@@ -13,7 +13,8 @@ import mtdna_classifier
13
  # --- IMPORTANT: UNCOMMENT AND CONFIGURE YOUR REAL API KEY ---
14
  import google.generativeai as genai
15
 
16
- genai.configure(api_key=os.getenv("GOOGLE_API_KEY"))
 
17
 
18
  import nltk
19
  from nltk.corpus import stopwords
@@ -22,12 +23,18 @@ try:
22
  except LookupError:
23
  nltk.download('stopwords')
24
  nltk.download('punkt_tab')
25
- # --- Define Pricing Constants (for Gemini 1.5 Flash & text-embedding-004) ---
26
- # Prices are per 1,000 tokens
27
- PRICE_PER_1K_INPUT_LLM = 0.000075 # $0.075 per 1M tokens
28
- PRICE_PER_1K_OUTPUT_LLM = 0.0003 # $0.30 per 1M tokens
29
- PRICE_PER_1K_EMBEDDING_INPUT = 0.000025 # $0.025 per 1M tokens
30
-
 
 
 
 
 
 
31
  # --- API Functions (REAL API FUNCTIONS) ---
32
 
33
  # def get_embedding(text, task_type="RETRIEVAL_DOCUMENT"):
@@ -59,7 +66,7 @@ def get_embedding(text, task_type="RETRIEVAL_DOCUMENT"):
59
  return np.zeros(768, dtype='float32')
60
 
61
 
62
- def call_llm_api(prompt, model_name='gemini-1.5-flash-latest'):
63
  """Calls a Google Gemini LLM with the given prompt."""
64
  try:
65
  model = genai.GenerativeModel(model_name)
@@ -1087,9 +1094,10 @@ def query_document_info(query_word, alternative_query_word, metadata, master_str
1087
  # 5. Execute RAG if needed (either full RAG or targeted RAG for missing fields)
1088
 
1089
  # Determine if a RAG call is necessary
1090
- run_rag = (extracted_country == 'unknown' or extracted_type == 'unknown')# or \
1091
- #extracted_ethnicity == 'unknown' or extracted_specific_location == 'unknown')
1092
- global_llm_model_for_counting_tokens = genai.GenerativeModel('gemini-1.5-flash-latest')
 
1093
  if run_rag:
1094
  print("try run rag")
1095
  # Determine the phrase for LLM query
 
13
  # --- IMPORTANT: UNCOMMENT AND CONFIGURE YOUR REAL API KEY ---
14
  import google.generativeai as genai
15
 
16
+ #genai.configure(api_key=os.getenv("GOOGLE_API_KEY"))
17
+ genai.configure(api_key=os.getenv("GOOGLE_API_KEY_BACKUP"))
18
 
19
  import nltk
20
  from nltk.corpus import stopwords
 
23
  except LookupError:
24
  nltk.download('stopwords')
25
  nltk.download('punkt_tab')
26
+ # # --- Define Pricing Constants (for Gemini 1.5 Flash & text-embedding-004) ---
27
+ # # Prices are per 1,000 tokens
28
+ # PRICE_PER_1K_INPUT_LLM = 0.000075 # $0.075 per 1M tokens
29
+ # PRICE_PER_1K_OUTPUT_LLM = 0.0003 # $0.30 per 1M tokens
30
+ # PRICE_PER_1K_EMBEDDING_INPUT = 0.000025 # $0.025 per 1M tokens
31
+
32
+ # Gemini 2.5 Flash-Lite pricing per 1,000 tokens
33
+ PRICE_PER_1K_INPUT_LLM = 0.00010 # $0.10 per 1M input tokens
34
+ PRICE_PER_1K_OUTPUT_LLM = 0.00040 # $0.40 per 1M output tokens
35
+
36
+ # Embedding-001 pricing per 1,000 input tokens
37
+ PRICE_PER_1K_EMBEDDING_INPUT = 0.00015 # $0.15 per 1M input tokens
38
  # --- API Functions (REAL API FUNCTIONS) ---
39
 
40
  # def get_embedding(text, task_type="RETRIEVAL_DOCUMENT"):
 
66
  return np.zeros(768, dtype='float32')
67
 
68
 
69
+ def call_llm_api(prompt, model_name="gemini-2.5-flash-lite"):#'gemini-1.5-flash-latest'):
70
  """Calls a Google Gemini LLM with the given prompt."""
71
  try:
72
  model = genai.GenerativeModel(model_name)
 
1094
  # 5. Execute RAG if needed (either full RAG or targeted RAG for missing fields)
1095
 
1096
  # Determine if a RAG call is necessary
1097
+ # run_rag = (extracted_country == 'unknown' or extracted_type == 'unknown')# or \
1098
+ # #extracted_ethnicity == 'unknown' or extracted_specific_location == 'unknown')
1099
+ run_rag = True
1100
+ global_llm_model_for_counting_tokens = genai.GenerativeModel("gemini-2.5-flash-lite")#('gemini-1.5-flash-latest')
1101
  if run_rag:
1102
  print("try run rag")
1103
  # Determine the phrase for LLM query