# Week - 3 Assignment - Integrate Traditional Chatbot with AI Service Project (Transformers) Praveen Kumar Parimi

#importing the required libraries including transformers
import base64
import gradio as gr
from huggingface_hub import InferenceClient
from transformers import AutoProcessor, AutoModelForImageTextToText
import torch

processor = AutoProcessor.from_pretrained("Qwen/Qwen2.5-VL-3B-Instruct")
model = AutoModelForImageTextToText.from_pretrained("Qwen/Qwen2.5-VL-3B-Instruct")

def chat(input_text, history=[]):
    history.append(input_text)
    prompt = "\n".join(history) + "\nAI:"  # Simple conversational format
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    output = model.generate(**inputs, max_length=512, pad_token_id=tokenizer.eos_token_id)
    response = tokenizer.decode(output[:, inputs.input_ids.shape[-1]:][0], skip_special_tokens=True)
    history.append(f"AI: {response}")
    return response, history

"""
For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
"""
print("starting Praveen's smarter chatbot...")

"""
The transformer model used here is Microsoft-trained Phi-3.5-mini-instruct
"""

#model_name = "microsoft/Phi-3.5-mini-instruct"

chat_model = pipeline("text-generation", model=model_name)

print("defining the chat_response function")

def chat_response(
    message,
    history: list[tuple[str, str]],
    system_message,
    max_tokens    
):

    print("Inside chat_response progressing...") 
    
    messages = [{"role": "system", "content": system_message}]

    print ("System Messages", messages)
    
    messages.append({"role": "user", "content": message})
    
    print ("Messages after adding user messages", messages)
    
    response = chat_model(messages)  #Passing system and user messages to the transformer model Phi-3.5-mini-instruct to get smarter responses
      
    print("Response received from model",response)
    
    return response[-1]['generated_text'][-1]['content']
"""
For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
"""

demo = gr.ChatInterface(
    chat,
    additional_inputs=[
        gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens")       
    ],
)


if __name__ == "__main__":
    demo.launch()