Spaces:
Sleeping
Sleeping
File size: 3,175 Bytes
440ae28 4210dea f5f674c 4210dea f5f674c 4210dea 440ae28 4210dea 2251617 4210dea dc6dfa5 4210dea dc6dfa5 4210dea f5f674c 4210dea f5f674c 4210dea 440ae28 4210dea | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 | from copy import copy
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain.output_parsers import PydanticOutputParser
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import PromptTemplate
from pydantic import BaseModel, Field
from dotenv import load_dotenv
from typing import Optional
import pandas as pd
import os
class InsuredDetails(BaseModel):
"""Insured business details"""
business_name: Optional[str] = Field(description="The insured business name")
business_address: str = Field(description="The insured business address")
contact_name: Optional[str] = Field(description="the person of contact's name")
contact_phone: Optional[str] = Field(description="the person of contact's phone number")
def __str__(self):
content = (f"business_name: {self.business_name}\nbusiness_address: {self.business_address}\n"
f"contact_name: {self.contact_name}\ncontact_phone: {self.contact_phone}")
separator = "~" * 20
return "\n".join([content, separator])
def parse_to_df(results):
return pd.DataFrame([ins.model_dump() for ins in results])
load_dotenv() # Only needed locally
google_api_key = os.getenv("GOOGLE_API_KEY")
if not google_api_key:
raise ValueError("GOOGLE_API_KEY not set. Please check your environment variables.")
def llm_parser(docs, model="gemini-2.0-flash-001"):
llm = ChatGoogleGenerativeAI(model=model, temperature=0)
parser = PydanticOutputParser(pydantic_object=InsuredDetails)
prompt = PromptTemplate(
template="""
Extract the insured business details from the following documents.
Make sure to extract the full address. If the address is spread across multiple lines
(e.g., "building number" and "street name" separately),
combine them to form a full address in the format: street name, building number, city.
The contact phone number format should be: 05xxxxxxxx (ten digits).
{format_instructions}
### Examples:
Text:
讻转讜讘转 诪讘讜讟讞: 9, 讟讘专讬讛
讛注专讜转 诇讻转讜讘转 讛诪讘讜讟讞: 转诪专 9
Output:
business_name: null
business_address: 转诪专 9, 讟讘专讬讛
contact_name: null
contact_phone: null
~~~~~~~~~~~~~~~~~~~~
Text:
Acme Corp is based at 123 Main St, Springfield. The contact person is Jane Doe, 055-123-4567.
Output:
business_name: Acme Corp
business_address: 123 Main St, Springfield
contact_name: Jane Doe
contact_phone: 0551234567
~~~~~~~~~~~~~~~~~~~~
### Now extract the details for the following document:
Text:
{input_text}
""",
input_variables=["input_text"],
partial_variables={"format_instructions": parser.get_format_instructions()},
)
chain = prompt | llm | StrOutputParser()
results = []
for doc in docs:
# Use the invoke method instead of run
raw_output = chain.invoke({"input_text": doc})
parsed = parser.parse(raw_output)
results.append(parsed)
for item in results:
print(item)
df_results = parse_to_df(results)
return results, df_results
|