Update app.py

app.py

@@ -1,45 +1,29 @@
-from langchain_huggingface import HuggingFacePipeline, ChatHuggingFace
-from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
-from profanity_check import predict
 import streamlit as st
-import re
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM
+from profanity_check import predict
+
 
-st.set_page_config(page_title="Chef Medi", layout="centered")
 st.title("Chef Medi")
 st.write("Share what you want to cook...")
 
-chat_template = ChatPromptTemplate.from_messages([
-    ('system', 'You are Medi, a concise 5-star Michelin chef AI. Teach in bullet points.'),
-    MessagesPlaceholder(variable_name='chat_history'),
-    ('human', '{user_input}')
-])
+prompt = st.chat_input("Say something")
 
 chat_history = []
 
-llm = HuggingFacePipeline.from_model_id(
-    model_id='TinyLlama/TinyLlama-1.1B-Chat-v1.0',
-    task='text-generation',
-    pipeline_kwargs=dict(
-        max_new_tokens=128,
-        temperature=0.7
-    )
-)
+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)
 
-model = ChatHuggingFace(llm=llm)
-
-def model_answer(user_input):
-    recent_history = chat_history[-3:]
-    recent_history.append({'role': 'user', 'content': user_input})
-    prompt = chat_template.invoke({
-        'chat_history': recent_history,
-        'user_input': user_input
-    })
-    result = model.invoke(prompt)
-    match = re.search(r"<\|assistant\|>(.*)", result.content, re.DOTALL)
-    ai_resp = match.group(1).strip()
-    return ai_resp
-
-prompt = st.chat_input("Say something")
+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
 
 if prompt:
     st.write(f"You: {prompt}")
@@ -47,7 +31,7 @@ if prompt:
     if offensive == [1]:
         ai_resp = "Please refrain from using bad language. Thanks"
     else:
-        ai_resp = model_answer(prompt)
+        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}")