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

# Load model and tokenizer
tokenizer = AutoTokenizer.from_pretrained("IEEEVITPune-AI-Team/ChatbotAlpha0.7")
model = AutoModelForCausalLM.from_pretrained("IEEEVITPune-AI-Team/ChatbotAlpha0.7")

# Define function to generate response
def generate_response(message, history, system_prompt, tokens):
    # Concatenate system prompt and user message
    input_text = f"{system_prompt} {message}"
    # Tokenize input text
    input_ids = tokenizer.encode(input_text, return_tensors="pt")
    # Generate response
    output = model.generate(input_ids, max_length=100, pad_token_id=tokenizer.eos_token_id)
    response = tokenizer.decode(output[0], skip_special_tokens=True)
    return response

# Define Gradio interface
with gr.Blocks() as demo:
    system_prompt = gr.Textbox("You are helpful AI.", label="System Prompt")
    slider = gr.Slider(10, 100, render=False, label="Number of Tokens")
    gr.ChatInterface(
        generate_response,
        inputs=["text", "text", system_prompt, slider],
        outputs="text"
    )

# Launch Gradio interface
demo.launch()