import gradio as gr
from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
import json
import torch

model_name = "shashankverma590/tiny-llama-1b-kid-friendly-chatbot-tiny"

tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16, device_map="auto")

pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)

# System prompt
system_prompt = """You are a helpful chatbot for conversing with kids under the age of 7.
You should be empathetic, encouraging, and positive-minded in general.
The current mood of the user is happy. You should reply accordingly."""

def get_chatbot_response(conversation_history_json):
    try:
        # Parse JSON input
        conversation_history = json.loads(conversation_history_json)
        print (conversation_history)
        
        # Create a formatted prompt
        # prompt = apply_chat_template(conversation_history)
        prompt = pipe.tokenizer.apply_chat_template(conversation_history, tokenize=False, add_generation_prompt=True)
        
        # Generate response
        outputs = pipe(prompt, max_new_tokens=256, do_sample=False, temperature=0.1, top_k=50, top_p=0.1, eos_token_id=pipe.tokenizer.eos_token_id, pad_token_id=pipe.tokenizer.pad_token_id)
        response = outputs[0]['generated_text'][len(prompt):].strip()
        return response
    except json.JSONDecodeError:
        return "Invalid input format. Please provide valid JSON conversation history."
    except Exception as e:
        return f"An error occurred: {str(e)}"

# Create the Gradio interface
app = gr.Interface(
    fn=get_chatbot_response,
    inputs=gr.Textbox(label="Your message (JSON format)", lines=5, placeholder='[{"user": "Hi!"}]'),
    outputs=gr.Textbox(label="System response"),
)

# Launch the app
app.launch(debug=True)