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

# Set the device to CPU since Hugging Face Spaces does not support GPU
device = torch.device("cpu")  # Ensure it's using CPU only

# Load model and tokenizer
model_name = "hosseinhimself/ISANG-v1.0-8B"  # Replace with your model name
tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=False)
model = AutoModelForCausalLM.from_pretrained(model_name).to(device)

# Define the Alpaca-style prompt template
alpaca_prompt = """Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.

### Instruction:
You are ISANG, a multilingual large language model made by ISANG AI. You only respond in Persian, Korean, or English. If a user uses one of these languages, reply in the same language.

### Input:
{}

### Response:
{}"""

# Function to generate responses
def generate_response(input_text, max_tokens=1024, temperature=0.7, history=[]):
    # Retain only the last two exchanges for context
    if len(history) > 2:
        history = history[-2:]
    
    # Format the prompt
    prompt = "\n".join(history + [f"User: {input_text}\nAI:"])
    
    # Tokenize the input
    inputs = tokenizer(prompt, return_tensors="pt").to(device)
    
    # Generate model output
    output = model.generate(
        inputs.input_ids,
        max_new_tokens=max_tokens,
        temperature=temperature
    )
    
    # Decode the model output
    response = tokenizer.decode(output[0], skip_special_tokens=True).strip()
    
    # Update the history
    history.append(f"User: {input_text}")
    history.append(f"AI: {response}")
    
    return response, history

# Gradio interface
iface = gr.Interface(
    fn=generate_response,
    inputs=[
        gr.Textbox(label="Your Message", placeholder="Type your message here..."),
        gr.Slider(minimum=1, maximum=2048, value=1024, step=1, label="Max Tokens"),
        gr.Slider(minimum=0.0, maximum=2.0, value=0.7, step=0.1, label="Temperature"),
        gr.State(value=[])  # To maintain conversation history
    ],
    outputs=[gr.Textbox(label="AI Response"), gr.State()],
    title="ISANG Chatbot",
    description="A chatbot powered by ISANG-v1.0-8B model. Chat with me!",
    theme="huggingface",  # Purple theme
    live=False  # Set to False since live updates aren't required
)

# Launch the interface
iface.launch()