-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathllm_utils.py
42 lines (36 loc) · 2.01 KB
/
llm_utils.py
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
from langchain.llms.ollama import Ollama
from langchain.schema import StrOutputParser
from langchain.prompts import PromptTemplate
from typing import List
class OllamaModel(object):
def __init__(self, ollama_model_name: str, temperature: float = 0) -> None:
"""
Initialize the Ollama Model.
Params:
ollama_model_name (str): The Ollama model name to use (recommended: mistral-openorca)
temperature (float): The temperature to use for the model (default: 0)
"""
self.ollama_model_name = ollama_model_name
self.temperature = temperature
self.llm = Ollama(model=ollama_model_name, temperature=temperature)
def run(self, items_list: List[str]) -> str:
"""Run the Model."""
write_template = PromptTemplate.from_template(
template="""
The description list is: "{inst}"
You are a helpful AI assistant designed to output a single phrase (2-3 words) and no more.
You are part of a code pipeline, and any output will be used downstream as written.
You must not output anything but your final answer, as other parts of the pipeline rely on a consistent output.
You will be a trusted and reliable part of this pipeline sticking to what you are told.
You will be given a list of movie descriptions, and you must output a phrase that all of them share by identifying similarities between them.
The category should be a short phrase describing the general theme of the movies.
Remember to only output your final phrase consisting of 2-3 words.
Your output should be short and concise.
Your final phrase is: """,
)
chain = write_template | self.llm | StrOutputParser()
output = chain.invoke({"inst": str(items_list)})
output = output.strip()
output = output.lower()
output = output[0].upper() + output[1:] # capitalize
return output