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

tokenizer = AutoTokenizer.from_pretrained("TinyLlama/TinyLlama-1.1B-intermediate-step-1431k-3T")
model = AutoModelForCausalLM.from_pretrained("TinyLlama/TinyLlama-1.1B-intermediate-step-1431k-3T")

def generate_text(system_prompt, user_prompt, temperature, max_length, min_length):
    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": user_prompt},
    ]
    tokenized_chat = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)

    inputs = tokenizer(tokenized_chat, return_tensors="pt").to(model.device) 

    outputs = model.generate(**inputs, use_cache=True, max_length=max_length, min_length=min_length, temperature=temperature, num_return_sequences=1)

    generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)

    return generated_text

def chatbot_app(system_prompt, user_prompt, temperature, max_length, min_length):
    generated_text = generate_text(system_prompt, user_prompt, temperature, max_length, min_length)
    return generated_text

iface = gr.Interface(
    fn=chatbot_app,
    inputs=[
        gr.Textbox(value="You are a helpful assistant. Your job is to convert given text into SQL query.", label="System Prompt"),
        gr.Textbox(value="how many employees are there in CS department?", label="User Prompt"),
        gr.Number(minimum=0.1, maximum=2.0, value=1.0, label="Temperature"),
        gr.Number(minimum=0, maximum=2048, value=128, label="Max Length"),
        gr.Number(minimum=0, maximum=2048, value=1, label="Min Length"),
    ],
    outputs="text",
    live=False,
)

iface.launch()