Spaces:
Runtime error
Runtime error
File size: 1,209 Bytes
1c63b9f a79e48c 904ef74 cd21998 28b4b1c a79e48c 28b4b1c 9c1f9ef a79e48c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 |
import os
import gradio as gr
from transformers import AutoTokenizer, pipeline
import torch
# Initialize the model and tokenizer
model_name = "AIFS/Prometh-MOEM-V.01"
HF_TOKEN = os.environ.get("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": torch.float16, "load_in_4bit": True},
use_auth_token=hf_token
)
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()
|