Trying / app.py
pokeberrypie's picture
Update app.py
9ffc3ef verified
import gradio as gr
from transformers import pipeline, AutoConfig
import logging
# Temporary workaround: Extend the Logger class to include warning_once
def warning_once(self, msg, *args, **kwargs):
if msg not in self._warned:
self.warning(msg, *args, **kwargs)
self._warned.add(msg)
logging.Logger.warning_once = warning_once
logging.Logger._warned = set()
# Load the model configuration with trust_remote_code to execute local configuration
config = AutoConfig.from_pretrained('tiiuae/falcon-40b-instruct', trust_remote_code=True)
# Load the model once when the script starts using the loaded config
generator = pipeline('text-generation', model='tiiuae/falcon-40b-instruct', config=config)
def generate_text(prompt):
# Use the preloaded model
return generator(prompt, max_length=100)[0]['generated_text']
def main():
with gr.Blocks() as demo:
gr.Markdown("## Text Generation Model")
gr.Markdown("This model generates text based on the input prompt. Powered by Hugging Face transformers.")
prompt = gr.Textbox(lines=2, placeholder="Type your prompt here...")
output = gr.Text(label="Generated Text")
prompt.change(fn=generate_text, inputs=prompt, outputs=output)
demo.launch()
if __name__ == "__main__":
main()