demo2.py

import pandas as pd
import json
from PIL import Image
import numpy as np

import os
from pathlib import Path

import torch
import torch.nn.functional as F

# from src.data.embs import ImageDataset
from src.model.blip_embs import blip_embs
from src.data.transforms import transform_test

from transformers import StoppingCriteria, StoppingCriteriaList, TextIteratorStreamer
import gradio as gr
# import spaces

from langchain.chains import ConversationChain
from langchain_community.chat_message_histories import ChatMessageHistory
from langchain_core.runnables import RunnableWithMessageHistory
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_groq import ChatGroq

from dotenv import load_dotenv


# GROQ_API_KEY = os.getenv("GROQ_API_KEY")
GROQ_API_KEY = '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'
load_dotenv(".env")
USER_AGENT = os.getenv("USER_AGENT")
GROQ_API_KEY = os.getenv("GROQ_API_KEY")
SECRET_KEY = os.getenv("SECRET_KEY")


# Set environment variables
os.environ['USER_AGENT'] = USER_AGENT
os.environ["GROQ_API_KEY"] = GROQ_API_KEY
os.environ["TOKENIZERS_PARALLELISM"] = 'true'

# Initialize LLM
llm = ChatGroq(model="llama-3.1-8b-instant", temperature=0, max_tokens=1024, max_retries=2)


# QA system prompt and chain
qa_system_prompt = """
Prompt:
You are a highly intelligent assistant. Use the following context to answer user questions. Analyze the data carefully and generate a clear, concise, and informative response to the user's question based on this data.

Response Guidelines:
    - Use only the information provided in the data to answer the question.
    - Ensure the answer is accurate and directly related to the question.
    - If the user question required no related data then give the truthful response to the user question.
    - If the data is insufficient to answer the question, politey apologise and tell the user that there is insufficient data available to answer their question.
    - Provide the response in a conversational yet friendly tone.

You are an AI assistant developed by Nutrigenics AI, specializing in intelligent recipe information retrieval. Your purpose is to help users by recommending recipes, providing detailed nutritional values, listing ingredients, offering step-by-step cooking instructions, and filtering recipes based on provide context ans user query.

Operational Guidelines:
1. Input Structure:
   - Context: You may receive contextual information related to recipes, such as specific data sets, user preferences, dietary restrictions, or previously selected dishes.
   - User Query: Users will pose questions or requests related to recipes, nutritional information, ingredient substitutions, cooking instructions, and more.

2. Response Strategy:
   - Utilize Provided Context: If the context contains relevant information that addresses the user's query, base your response on this provided data to ensure accuracy and relevance.
   - Respond to User Query Directly: If the context does not contain the necessary information to answer the user's query, generate a response based solely on the user's input and your trained knowledge.

Core Functionalities:
- Nutritional Information: Accurately provide nutritional values for each recipe, including calories, macronutrients (proteins, fats, carbohydrates), and essential vitamins and minerals, using contextual data when available.
- Ingredient Details: List all ingredients required for recipes, including substitute options for dietary restrictions or ingredient availability, utilizing context when relevant.
- Step-by-Step Cooking Instructions: Deliver clear, easy-to-follow instructions for preparing and cooking meals, informed by any provided contextual data.
- Recipe Recommendations: Suggest dishes based on user preferences, dietary restrictions, available ingredients, and contextual data if provided.

Additional Instructions:
- Precision and Personalization: Always aim to provide precise, personalized, and relevant information to users based on both the provided context and their specific queries.
- Clarity and Coherence: Ensure that all responses are clear, well-structured, and easy to understand, facilitating a seamless user experience.
- Substitute Suggestions: When suggesting ingredient substitutes, consider user preferences and dietary restrictions outlined in the context or user query.
- Dynamic Adaptation: Adapt your responses dynamically based on whether the context is relevant to the user's current request, ensuring optimal use of available information.

Don't mention about context in the response, format the answer in a natural friedly way.

Context:
{context}
"""
qa_prompt = ChatPromptTemplate.from_messages(
    [
        ("system", qa_system_prompt),
        ("human", "{input}")
    ]
)

# Create the base chain
base_chain = qa_prompt | llm | StrOutputParser()

# Wrap the chain with message history
question_answer_chain = RunnableWithMessageHistory(
    base_chain,
    lambda session_id: ChatMessageHistory(),  # This creates a new history for each session
    input_messages_key="input",
    history_messages_key="chat_history"
)


class StoppingCriteriaSub(StoppingCriteria):

    def __init__(self, stops=[], encounters=1):
        super().__init__()
        self.stops = stops

    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor):
        for stop in self.stops:
            if torch.all(input_ids[:, -len(stop):] == stop).item():
                return True

        return False

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

def get_blip_config(model="base"):
    config = dict()
    if model == "base":
        config[
            "pretrained"
        ] = "https://storage.googleapis.com/sfr-vision-language-research/BLIP/models/model_base_capfilt_large.pth "
        config["vit"] = "base"
        config["batch_size_train"] = 32
        config["batch_size_test"] = 16
        config["vit_grad_ckpt"] = True
        config["vit_ckpt_layer"] = 4
        config["init_lr"] = 1e-5
    elif model == "large":
        config[
            "pretrained"
        ] = "https://storage.googleapis.com/sfr-vision-language-research/BLIP/models/model_large_retrieval_coco.pth"
        config["vit"] = "large"
        config["batch_size_train"] = 16
        config["batch_size_test"] = 32
        config["vit_grad_ckpt"] = True
        config["vit_ckpt_layer"] = 12
        config["init_lr"] = 5e-6

    config["image_size"] = 384
    config["queue_size"] = 57600
    config["alpha"] = 0.4
    config["k_test"] = 256
    config["negative_all_rank"] = True

    return config


print("Creating model")
config = get_blip_config("large")

model = blip_embs(
        pretrained=config["pretrained"],
        image_size=config["image_size"],
        vit=config["vit"],
        vit_grad_ckpt=config["vit_grad_ckpt"],
        vit_ckpt_layer=config["vit_ckpt_layer"],
        queue_size=config["queue_size"],
        negative_all_rank=config["negative_all_rank"],
    )

model = model.to(device)
model.eval()
print("Model Loaded !")
print("="*50)

transform = transform_test(384)

print("Loading Data")
df = pd.read_json("datasets/sidechef/my_recipes.json")

print("Loading Target Embedding")
tar_img_feats = []
for _id in df["id_"].tolist():     
    tar_img_feats.append(torch.load("datasets/sidechef/blip-embs-large/{:07d}.pth".format(_id)).unsqueeze(0))

tar_img_feats = torch.cat(tar_img_feats, dim=0)


class Chat:

    def __init__(self, model, transform, dataframe, tar_img_feats, device='cuda:0', stopping_criteria=None):
        self.device = device
        self.model = model
        self.transform = transform
        self.df = dataframe
        self.tar_img_feats = tar_img_feats
        self.img_feats = None
        self.target_recipe = None
        self.messages = []

        if stopping_criteria is not None:
            self.stopping_criteria = stopping_criteria
        else:
            stop_words_ids = [torch.tensor([2]).to(self.device)]
            self.stopping_criteria = StoppingCriteriaList([StoppingCriteriaSub(stops=stop_words_ids)])

    def encode_image(self, image_path):
        img = Image.fromarray(image_path).convert("RGB")
        img = self.transform(img).unsqueeze(0)
        img = img.to(self.device)
        img_embs = model.visual_encoder(img)
        img_feats = F.normalize(model.vision_proj(img_embs[:, 0, :]), dim=-1).cpu()

        self.img_feats = img_feats 

        self.get_target(self.img_feats, self.tar_img_feats)

    def get_target(self, img_feats, tar_img_feats) : 
        score = (img_feats @ tar_img_feats.t()).squeeze(0).cpu().detach().numpy()
        index = np.argsort(score)[::-1][0]
        self.target_recipe = df.iloc[index]

    def ask(self):
        return json.dumps(self.target_recipe.to_json())



chat = Chat(model,transform,df,tar_img_feats, device)
print("Chat Initialized !")


custom_css = """
.primary{
    background-color: #4CAF50; /* Green */
}
"""


# @spaces.GPU
def respond_to_user(image, message):
    # Process the image and message here
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    chat = Chat(model,transform,df,tar_img_feats, device)
    chat.encode_image(image)
    data = chat.ask()
    formated_input = {
        'input': message,
        'context': data
    }
    try:
        response = question_answer_chain.invoke(formated_input, config={"configurable": {"session_id": 'abc123'}})
    except Exception as e:
        print(e)
        response = {'content':"An error occurred while processing your request."}
    return response, data

iface = gr.Interface(
    fn=respond_to_user,
    inputs=[gr.Image(), gr.Textbox(label="Ask Query")],
    outputs=[gr.Textbox(label="Nutrition-GPT"), gr.JSON(label="context")],
    title="Nutrition-GPT Demo",
    description="Upload an food image and ask queries!",
    css=".component-12 {background-color: red}",
)

iface.launch()