from transformers import AutoModelForCausalLM, AutoTokenizer, FalconForCausalLM
fine_tuned_model = "ashioyajotham/falcon-coder"
model = FalconForCausalLM.from_pretrained(fine_tuned_model)
# Load the tokenizer for the Falcon 7B model with remote code trust
tokenizer = AutoTokenizer.from_pretrained(fine_tuned_model, trust_remote_code=True)

# Set the padding token to be the same as the end-of-sequence token
tokenizer.pad_token = tokenizer.eos_token

prompt = "Generate a Python script to add prime numbers between one and ten"

inputs = tokenizer.encode(prompt, return_tensors='pt')

outputs = model.generate(inputs, max_length=100, temperature = .7, do_sample=True, pad_token_id=tokenizer.eos_token_id)

completion = tokenizer.decode(outputs[0])

#print(completion)
import gradio as gr

iface = gr.Interface(
    fn=completion,
    inputs="text",
    outputs="text",
    title="Code Generation App",
    description="Generate code from text input."
)

iface.launch()