Update app.py

app.py

@@ -1,37 +1,55 @@
+from langchain_huggingface import HuggingFacePipeline, ChatHuggingFace
+from langchain_core.prompts import ChatPromptTemplate
 import streamlit as st
-import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM
 from profanity_check import predict
 
 
 st.title("Chef Medi")
 st.write("Share what you want to cook...")
 
-prompt = st.chat_input("Say something")
 
-chat_history = []
+# Chat Template
+
+chat_template = ChatPromptTemplate.from_messages([
+    ('system', 
+     'Your name is Medi and you are an AI five star michelin star chef who teaches cooking. Your job is to guide users and help them create delicious food. Write minimal text to teach users answer in bullet points'),
+    ('human', '{user_input}')
+])
+
+
+# Model
+
+llm = HuggingFacePipeline.from_model_id(
+    model_id='TinyLlama/TinyLlama-1.1B-Chat-v1.0',
+    task='text-generation',
+    pipeline_kwargs=dict(
+        max_new_tokens = 512
+    )
+)
+
+model = ChatHuggingFace(llm=llm)
 
-model_id = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
-tokenizer = AutoTokenizer.from_pretrained(model_id)
-model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.float16)
 
-def generate_reply(prompt):
-    chat_history.append(f"User: {prompt}")
-    prompt = "You are Medi, a 5-star chef. Be brief and answer in bullet points.\n" + "\n".join(chat_history) + "\nAssistant:"
-    inputs = tokenizer(prompt, return_tensors="pt")
-    outputs = model.generate(**inputs, max_new_tokens=1024, do_sample=False)
-    full_output = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    reply = full_output.split("Assistant:")[-1].strip()
-    chat_history.append(f"Assistant: {reply}")
-    return reply
+def model_answer(user_input, chat_template):
+
+    chain = chat_template | model
+
+    result = chain.invoke({
+        'user_input': user_input
+    })
+    
+    match = re.search(r"<\|assistant\|>(.*)", result.content, re.DOTALL)
+    ai_resp = match.group(1).strip()
+    return ai_resp
+
+
+prompt = st.chat_input("Say something")
 
 if prompt:
     st.write(f"You: {prompt}")
     offensive = predict([prompt])
     if offensive == [1]:
-        ai_resp = "Please refrain from using bad language. Thanks"
+        ai_resp = "Please refrain from use bad language. Thanks"
     else:
-        ai_resp = generate_reply(prompt)
-        chat_history.append({'role': 'user', 'content': prompt})
-        chat_history.append({'role': 'assistant', 'content': ai_resp})
-    st.write(f"Medi: {ai_resp}")
+        ai_resp = model_answer(prompt, chat_template)
+    st.write(f"Medi: {ai_resp}")