Trying / app.py
pokeberrypie's picture
Update app.py
648aefb verified
raw
history blame
905 Bytes
# app.py
import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer
def main():
# Load the model and tokenizer
model_name = "/tiiuae/falcon-180B" # Replace with the correct model name if different
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)
# Define the function that uses the model
def generate_text(input_text):
inputs = tokenizer(input_text, return_tensors="pt", max_length=512, truncation=True)
outputs = model.generate(**inputs, max_length=512, min_length=50, num_return_sequences=1)
return tokenizer.decode(outputs[0], skip_special_tokens=True)
# Create and launch the Gradio interface
iface = gr.Interface(fn=generate_text, inputs="text", outputs="text", title="Falcon 180B Text Generator")
iface.launch()
if __name__ == "__main__":
main()