Upload hogwats_gemini.py

hogwats_gemini.py

@@ -0,0 +1,65 @@
+import google.generativeai as genai
+from config import gemini_api
+import tiktoken
+
+def get_answer(query, company_name, chunked_raw_content):
+
+    genai.configure(api_key=gemini_api)
+
+    # Create the model
+    # See https://ai.google.dev/api/python/google/generativeai/GenerativeModel
+    generation_config = {
+    "temperature": 1,
+    "top_p": 0.95,
+    "top_k": 64,
+    "max_output_tokens": 8192,
+    "response_mime_type": "text/plain",
+    }
+    safety_settings = [
+    {
+        "category": "HARM_CATEGORY_HARASSMENT",
+        "threshold": "BLOCK_MEDIUM_AND_ABOVE",
+    },
+    {
+        "category": "HARM_CATEGORY_HATE_SPEECH",
+        "threshold": "BLOCK_MEDIUM_AND_ABOVE",
+    },
+    {
+        "category": "HARM_CATEGORY_SEXUALLY_EXPLICIT",
+        "threshold": "BLOCK_MEDIUM_AND_ABOVE",
+    },
+    {
+        "category": "HARM_CATEGORY_DANGEROUS_CONTENT",
+        "threshold": "BLOCK_MEDIUM_AND_ABOVE",
+    },
+    ]
+    chunks = []
+    for chunk in chunked_raw_content:
+        chunk = chunk.replace("PDF_FILE_____data_dumpster_", "")
+        chunk = chunk.replace("data_dumpster_", "")
+        chunks.append(chunk)
+    context = str(chunks)
+        
+    enc = tiktoken.get_encoding("cl100k_base")
+
+    
+    toks = enc.encode(context)
+    if len(toks) >= 900000:
+        chunk_size = int(len(context)  // (len(toks) / 900000))
+        chunk_size
+        context = context[:chunk_size]
+
+
+    model = genai.GenerativeModel(
+    model_name="gemini-1.5-flash-latest",
+    safety_settings=safety_settings,
+    generation_config=generation_config,
+    system_instruction=f"You are an expert at a Private Equity fund. You are helping a colleague with his due diligence on {company_name}. All the questions you will receive are in the context of this due diligence. You always precise the sources from the context (given below) you use.\nYou answer any question based on the following context elements:\n{context}",
+    )
+    chat_session = model.start_chat(
+    history=[],
+    )
+
+
+    response = chat_session.send_message(f"{query} - (Bain style answer + sources properly renamed if needed)")
+    return response.text