File size: 905 Bytes
caa0c26
7c2904b
 
 
 
 
 
648aefb
7c2904b
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
# 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()