Spaces:
Running
Running
Update model.py
Browse files
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 |
-
|
1092 |
-
|
|
|
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
|