langchain-ollama-writer.py

from langchain_community.llms import Ollama

from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser

import time

MODEL = "phi3"
SYSTEM_PROMPT = "You are a world class technical documentation writer."

prompt = ChatPromptTemplate.from_messages([
    ("system", SYSTEM_PROMPT),
    ("user", "{input}")
])

print(f'Ollama is loading {MODEL} model...')
llm = Ollama(model=MODEL)

output_parser = StrOutputParser()

chain = prompt | llm | output_parser

# loop to ask questions
while True:
    # Read text from stdin
    text = input("Enter question: ")

    print(f'Asking {MODEL} a question...')

    # Run query and messaure time
    start = time.time()

    result = chain.invoke({"input": text})

    end = time.time()

    # print time round to two decimal places
    print("Time taken: {:.2f} seconds".format(end - start))
    print(f'{MODEL} says: {result}')

    # Ask if user wants to ask another question
    another = input("Do you want to ask another question? (y/n): ")
    if another.lower() != "y":
        break