BramVanroy / app.py
rheremans's picture
Update app.py
95662cd verified
from transformers import pipeline, Conversation, AutoTokenizer, AutoModelForCausalLM
from langchain.llms import HuggingFacePipeline
from langchain.prompts import PromptTemplate
from langchain_community.llms import HuggingFaceHub
from langchain.callbacks.manager import CallbackManager
from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
from langchain.chains import LLMChain
#1: "meta-llama/Llama-2-13b-chat-hf",
#2: "BramVanroy/Llama-2-13b-chat-dutch"
my_config = {'model_name': "meta-llama/Llama-2-13b-chat-hf", #"./Bram", #BramVanroy/Llama-2-13b-chat-dutch",
'do_sample': True, 'temperature': 0.1,
'repetition_penalty': 1.1, 'max_new_tokens': 500, }
print(f"Selected model: {my_config['model_name']}")
print(f"Parameters are: {my_config}")
question = "Who won the FIFA World Cup in the year 1994? "
template = """Question: {question}
Answer: Let's think step by step."""
prompt = PromptTemplate.from_template(template)
def generate_with_llama_chat(my_config):
print('tokenizer')
tokenizer = AutoTokenizer.from_pretrained(my_config['model_name'])
print('causal')
model = AutoModelForCausalLM.from_pretrained(my_config['model_name'])
print('Pipeline')
chatbot = pipeline("text-generation",model=my_config['model_name'],
tokenizer=tokenizer,
do_sample=my_config['do_sample'],
temperature=my_config['temperature'],
repetition_penalty=my_config['repetition_penalty'],
#max_length=my_config['max_length'],
max_new_tokens=my_config['max_new_tokens'],
model_kwargs={"device_map": "auto","load_in_8bit": True})
return chatbot
llama_chat = generate_with_llama_chat(my_config)
# Set up callback manager to print output word by word
callback_manager = CallbackManager([StreamingStdOutCallbackHandler()])
llm = HuggingFacePipeline(pipeline=llama_chat, callback_manager=callback_manager)
llm_chain = LLMChain(prompt=prompt, llm=llm)
print(llm_chain.invoke(question))