import os
import gradio as gr
from transformers import AutoTokenizer, pipeline

# Initialize the model and tokenizer
model_name = "AIFS/Prometh-MOEM-V.01"
hf_token = os.getenv("HF_TOKEN")

tokenizer = AutoTokenizer.from_pretrained(model_name, use_auth_token=hf_token)
text_generation_pipeline = pipeline(
    "text-generation",
    model=model_name,
    model_kwargs={"torch_dtype": "auto", "load_in_4bit": True},
)

def generate_text(user_input):
    messages = [{"role": "user", "content": user_input}]
    prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    outputs = text_generation_pipeline(prompt, max_new_tokens=256, do_sample=True, temperature=0.7, top_k=50, top_p=0.95)
    return outputs[0]["generated_text"]

# Create the Gradio interface
iface = gr.Interface(
    fn=generate_text,
    inputs=gr.Textbox(lines=2, placeholder="Type your question here..."),
    outputs=gr.Textbox(),
    title="Prometh-MOEM Text Generation",
    description="A text generation model that understands your queries and generates concise, informative responses."
)

# Launch the interface (omit `share=True` when deploying on Hugging Face Spaces)
iface.launch()