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

logging.basicConfig(level=logging.DEBUG, format="%(asctime)s - %(levelname)s - %(message)s")

model_name = "deepseek-ai/DeepSeek-R1"
logging.info(f"Loading model: {model_name}")
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)

def wrapped_model(input_text):
    logging.info(f"Input text: {input_text}")

    inputs = tokenizer(input_text, return_tensors="pt")
    logging.debug(f"Tokenized input: {inputs}")

    outputs = model.generate(**inputs, max_length=512, num_return_sequences=1)
    logging.debug(f"Model output tokens: {outputs}")

    decoded_output = tokenizer.decode(outputs[0], skip_special_tokens=True)
    logging.info(f"Decoded output: {decoded_output}")

    return decoded_output

interface = gr.Interface(
    fn=wrapped_model,
    inputs=gr.Textbox(lines=2, label="Enter your prompt"),
    outputs=gr.Textbox(label="Output")
)

interface.launch()