image.py 1.7 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253
  1. from langchain_community.llms import Ollama
  2. from langchain_community.embeddings import OllamaEmbeddings
  3. from langchain.prompts import ChatPromptTemplate
  4. from langchain_core.output_parsers import StrOutputParser
  5. from langchain_community.document_loaders.image import UnstructuredImageLoader
  6. from langchain.chains.combine_documents import create_stuff_documents_chain
  7. from langchain_community.vectorstores import Chroma
  8. from langchain.chains import create_retrieval_chain
  9. from langchain.text_splitter import RecursiveCharacterTextSplitter
  10. MODEL = "llava"
  11. prompt = ChatPromptTemplate.from_template(
  12. """
  13. You are a world class expert in vehicle appraisal.
  14. Answer only based on the following provided context. If you know the answer but it's not based in the
  15. provided context, don't provide the answer, just state the answer is not in the context provided:
  16. <context>
  17. {context}
  18. </context>
  19. User: {input}
  20. """
  21. )
  22. llm = Ollama(model=MODEL)
  23. loader = UnstructuredImageLoader("img/screenshot.jpg")
  24. data = loader.load()
  25. text_splitter = RecursiveCharacterTextSplitter(chunk_size=1024, chunk_overlap=100)
  26. chunks = text_splitter.split_documents(data)
  27. vector = Chroma.from_documents(data, OllamaEmbeddings())
  28. retriever = vector.as_retriever()
  29. output_parser = StrOutputParser()
  30. document_chain = create_stuff_documents_chain(llm, prompt, output_parser=output_parser)
  31. retrieval_chain = create_retrieval_chain(retriever, document_chain)
  32. response = retrieval_chain.invoke({
  33. "input": """
  34. Estás viendo la publicación de un vehículo, búsca la siguiente información:
  35. año
  36. marca
  37. modelo
  38. precio
  39. kilometraje
  40. combustible
  41. transmisión
  42. código publicación
  43. """
  44. })
  45. print(response["answer"])