ai-sandbox/scrape/image.py

from langchain_community.llms import Ollama
from langchain_community.embeddings import OllamaEmbeddings
from langchain.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
from langchain_community.document_loaders.image import UnstructuredImageLoader
from langchain.chains.combine_documents import create_stuff_documents_chain
from langchain_community.vectorstores import Chroma
from langchain.chains import create_retrieval_chain
from langchain.text_splitter import RecursiveCharacterTextSplitter

MODEL = "llava"

prompt = ChatPromptTemplate.from_template(
    """
    You are a world class expert in vehicle appraisal.
    Answer only based on the following provided context. If you know the answer but it's not based in the 
    provided context, don't provide the answer, just state the answer is not in the context provided:
    <context>
    {context}
    </context>

    User: {input}
    """
)

llm = Ollama(model=MODEL)

loader = UnstructuredImageLoader("img/screenshot.jpg")
data = loader.load()
text_splitter = RecursiveCharacterTextSplitter(chunk_size=1024, chunk_overlap=100)
chunks = text_splitter.split_documents(data)
vector = Chroma.from_documents(data, OllamaEmbeddings())
retriever = vector.as_retriever()
output_parser = StrOutputParser()

document_chain = create_stuff_documents_chain(llm, prompt, output_parser=output_parser)
retrieval_chain = create_retrieval_chain(retriever, document_chain)

response = retrieval_chain.invoke({
    "input": """
    Estás viendo la publicación de un vehículo, búsca la siguiente información:
    año
    marca
    modelo
    precio
    kilometraje
    combustible
    transmisión
    código publicación
    """
})
print(response["answer"])