commit

langchain_pipeline.py

@@ -2,8 +2,8 @@ import os
 from pdfminer import high_level
 from langchain_core.prompts import PromptTemplate
 from google.cloud import aiplatform
-from google.cloud.aiplatform import ModelServiceClient
-from google.cloud.aiplatform.gapic.schema import predict
+from google.cloud.aiplatform_v1 import ModelServiceClient
+from google.cloud.aiplatform_v1.types import GenerateContentRequest, Document, GenerationConfig, SafetySettings, HarmCategory, HarmBlockThreshold
 import base64
 
 # Initialize the Google AI Platform
@@ -11,15 +11,15 @@ aiplatform.init(project="akroda", location="us-central1")
 
 # Define the documents (base64-encoded PDF content)
 documents = [
-    predict.instance.TextDocument(content_type="application/pdf", data=base64.b64decode("JVBERi0xLjQKJeODgxNz5dL1Jvb3QgMTU0IDAgUi9TaXplIDE2Nj4+CnN0YXJ0eHJlZgoyMTY0NjkKJSVFT0YK")),
-    predict.instance.TextDocument(content_type="application/pdf", data=base64.b64decode("JVBERi0xLjQKJeLjz9MKNijU+PgpzdGFydHhyZWYKMTMxMDY0CiUlRU9GCg==")),
-    predict.instance.TextDocument(content_type="application/pdf", data=base64.b64decode("JVBERi0xLjQKJeLjz9MKNiAwZDU0YTVlNzllMWRhYWY1ZDQ2YjI+XS9Sb290IDE3NyAwIFIvU2l6ZSAxODc+PgpzdGFydHhyZWYKMjA3NTk5CiUlRU9GCg==")),
-    predict.instance.TextDocument(content_type="application/pdf", data=base64.b64decode("JVBERi0xLjQKJeLjz9ML1Jvb3QgMTg5IDAgUi9TaXplIDE5OT4+CnN0YXJ0eHJlZgoxOTgzNzMKJSVFT0YK")),
-    predict.instance.TextDocument(content_type="application/pdf", data=base64.b64decode("JVBERi0xCcnCmVuZHN0cmVhbQplbmRvYmoKc3RhcnR4cmVmCjIwOTgyNQolJUVPRgo=")),
-    predict.instance.TextDocument(content_type="application/pdf", data=base64.b64decode("JVBERi0xLj+CnN0YXJ0eHJlZgoyMTk5MDYKJSVFT0YK")),
-    predict.instance.TextDocument(content_type="application/pdf", data=base64.b64decode("JVBERi0xLjQKJiUlRU9GCg==")),
-    predict.instance.TextDocument(content_type="application/pdf", data=base64.b64decode("JVBERi0xLjQKJe90IDMwOCAwIFIvU2l6ZSAzMTg+PgpzdGFydHhyZWYKMjcwNzU3CiUlRU9GCg==")),
-    predict.instance.TextDocument(content_type="application/pdf", data=base64.b64decode("JVBERi0xLjUNJeLjz9MNCjcgMCBvYmoNPDwvTGluZWFyaXplZCAxL0wgNjc1NzgvTyA5L0UgNjAyNDYvTiAxL1QgNjcyODcvSCBbIDQ4MyAxNTRdPj4NZW5kb2JxDSAgICAgICAgICAgICAgICAgICAgDQoyMiAwIG9iag08PC9EZWNvZGVQYXJtczw8L0NvbHVtbnMgNC9QcmVkaWN0b3IgMTI+Pi9GaWx0ZXIvRmxhdGVEZWNvZGUvSURbPDE3NzU4MkJFODc4MzRFQjNBOEM3RkIzQTgyRjFFMEFCPjw5MzI2Qjk4REM4NjQ2RTRCODI3MzZFQUEzOENEQjFBQj5dL0luZGV4WzcgMjhdL0luZm8gNiAwIFIvTGVuZ3RoIDgzL1ByZXYgNjcyODgvUm9vdCA4IDAgUi9TaXplIDM1L1R5cGUvWFPRg0K"))
+    Document(content_type="application/pdf", data=base64.b64decode("JVBERi0xLjQKJeODgxNz5dL1Jvb3QgMTU0IDAgUi9TaXplIDE2Nj4+CnN0YXJ0eHJlZgoyMTY0NjkKJSVFT0YK")),
+    Document(content_type="application/pdf", data=base64.b64decode("JVBERi0xLjQKJeLjz9MKNijU+PgpzdGFydHhyZWYKMTMxMDY0CiUlRU9GCg==")),
+    Document(content_type="application/pdf", data=base64.b64decode("JVBERi0xLjQKJeLjz9MKNiAwZDU0YTVlNzllMWRhYWY1ZDQ2YjI+XS9Sb290IDE3NyAwIFIvU2l6ZSAxODc+PgpzdGFydHhyZWYKMjA3NTk5CiUlRU9GCg==")),
+    Document(content_type="application/pdf", data=base64.b64decode("JVBERi0xLjQKJeLjz9ML1Jvb3QgMTg5IDAgUi9TaXplIDE5OT4+CnN0YXJ0eHJlZgoxOTgzNzMKJSVFT0YK")),
+    Document(content_type="application/pdf", data=base64.b64decode("JVBERi0xCcnCmVuZHN0cmVhbQplbmRvYmoKc3RhcnR4cmVmCjIwOTgyNQolJUVPRgo=")),
+    Document(content_type="application/pdf", data=base64.b64decode("JVBERi0xLj+CnN0YXJ0eHJlZgoyMTk5MDYKJSVFT0YK")),
+    Document(content_type="application/pdf", data=base64.b64decode("JVBERi0xLjQKJiUlRU9GCg==")),
+    Document(content_type="application/pdf", data=base64.b64decode("JVBERi0xLjQKJe90IDMwOCAwIFIvU2l6ZSAzMTg+PgpzdGFydHhyZWYKMjcwNzU3CiUlRU9GCg==")),
+    Document(content_type="application/pdf", data=base64.b64decode("JVBERi0xLjUNJeLjz9MNCjcgMCBvYmoNPDwvTGluZWFyaXplZCAxL0wgNjc1NzgvTyA5L0UgNjAyNDYvTiAxL1QgNjcyODcvSCBbIDQ4MyAxNTRdPj4NZW5kb2JxDSAgICAgICAgICAgICAgICAgICAgDQoyMiAwIG9iag08PC9EZWNvZGVQYXJtczw8L0NvbHVtbnMgNC9QcmVkaWN0b3IgMTI+Pi9GaWx0ZXIvRmxhdGVEZWNvZGUvSURbPDE3NzU4MkJFODc4MzRFQjNBOEM3RkIzQTgyRjFFMEFCPjw5MzI2Qjk4REM4NjQ2RTRCODI3MzZFQUEzOENEQjFBQj5dL0luZGV4WzcgMjhdL0luZm8gNiAwIFIvTGVuZ3RoIDgzL1ByZXYgNjcyODgvUm9vdCA4IDAgUi9TaXplIDM1L1R5cGUvWFPRg0K"))
 ]
 
 text1 = """
@@ -64,18 +64,18 @@ val = prompt.format(
     min_transaction_overdraft=min_transaction_overdraft,
 )
 
-generation_config = generative_model.GenerationConfig(
+generation_config = GenerationConfig(
     max_output_tokens=8192,
     temperature=1,
     top_p=0.95,
 )
 
-safety_settings = generative_model.SafetySettings(
+safety_settings = SafetySettings(
     harm_category_settings={
-        generative_model.HarmCategory.HARM_CATEGORY_HATE_SPEECH: generative_model.HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE,
-        generative_model.HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT: generative_model.HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE,
-        generative_model.HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT: generative_model.HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE,
-        generative_model.HarmCategory.HARM_CATEGORY_HARASSMENT: generative_model.HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE,
+        HarmCategory.HARM_CATEGORY_HATE_SPEECH: HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE,
+        HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT: HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE,
+        HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT: HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE,
+        HarmCategory.HARM_CATEGORY_HARASSMENT: HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE,
     }
 )
 
@@ -83,7 +83,7 @@ def generate(document_parts, prompt_text):
     model_service_client = ModelServiceClient()
     model_resource_name = model_service_client.model_path("akroda", "us-central1", "gemini-1.5-pro-001")
     response = model_service_client.generate_content(
-        request=generative_model.GenerateContentRequest(
+        request=GenerateContentRequest(
             model=model_resource_name,
             documents=document_parts,
             prompt=prompt_text,