wizardlm_api / app.py
DR-Rakshitha's picture
Update app.py
e5c60ed
raw
history blame
1.03 kB
import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
# Specify the directory containing the tokenizer's configuration file (config.json)
model_name = "pytorch_model-00001-of-00002.bin"
# Initialize the tokenizer
# tokenizer = AutoTokenizer.from_pretrained(model_name, local_files_only=True)
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
tokenizer.pad_token = tokenizer.eos_token
tokenizer.padding_side = "right"
# Initialize the GPT4All model
model = AutoModelForCausalLM.from_pretrained(model_name)
def generate_text(input_text):
pipe = pipeline(task="text-generation", model=model, tokenizer=tokenizer, max_length=200)
result = pipe(f"<s>[INST] {input_text} [/INST]")
return result[0]['generated_text']
text_generation_interface = gr.Interface(
fn=generate_text,
inputs=[
gr.inputs.Textbox(label="Input Text"),
],
outputs=gr.outputs.Textbox(label="Generated Text"),
title="GPT-4 Text Generation",
).launch()