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

# Initialize the model and tokenizer
model_name = "Wtzwho/Prometh-MOEM-V.01"
tokenizer = AutoTokenizer.from_pretrained(model_name)
text_generation_pipeline = pipeline(
    "text-generation",
    model=model_name,
    model_kwargs={"torch_dtype": torch.float16, "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.inputs.Textbox(lines=2, placeholder="Type your question here..."),
    outputs=gr.outputs.Textbox(),
    title="Prometh-MOEM Text Generation",
    description="A text generation model that understands your queries and generates concise, informative responses."
)

# Run the interface
iface.launch()