|
import os |
|
os.system("pip install torch sentencepiece transformers Xformers accelerate") |
|
import gradio as gr |
|
import torch |
|
from transformers import AutoModelForCausalLM, AutoTokenizer |
|
|
|
|
|
model = AutoModelForCausalLM.from_pretrained("cyberagent/open-calm-1b", device_map="auto", torch_dtype=torch.float16) |
|
tokenizer = AutoTokenizer.from_pretrained("cyberagent/open-calm-1b") |
|
|
|
|
|
def generate_text(input_text): |
|
inputs = tokenizer(input_text, return_tensors="pt").to(model.device) |
|
with torch.no_grad(): |
|
tokens = model.generate( |
|
**inputs, |
|
max_new_tokens=64, |
|
do_sample=True, |
|
temperature=0.7, |
|
top_p=0.9, |
|
repetition_penalty=1.05, |
|
pad_token_id=tokenizer.pad_token_id, |
|
) |
|
output = tokenizer.decode(tokens[0], skip_special_tokens=True) |
|
return output |
|
|
|
|
|
input_text = gr.inputs.Textbox(lines=2, label="入力テキスト") |
|
output_text = gr.outputs.Textbox(label="生成されたテキスト") |
|
|
|
|
|
iface = gr.Interface(fn=generate_text, inputs=input_text, outputs=output_text) |
|
|
|
|
|
iface.launch() |
|
|
|
|