Spaces:
Sleeping
Sleeping
| 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 | |