Spaces:
Runtime error
Runtime error
import gradio as gr | |
from transformers import AutoTokenizer, AutoModelForCausalLM | |
import torch | |
# ูู ุงุฐุฌ ุฃููุฏุฉ ูู ูุชูุญุฉ | |
models = { | |
"CodeGen 2B": "Salesforce/codegen-2B-multi", | |
"CodeParrot": "codeparrot/codeparrot-small", | |
"GPT-J-6B": "EleutherAI/gpt-j-6B", | |
"GPT2": "gpt2" # ูู ูุฐุฌ ุจุณูุท ูู fallback | |
} | |
# ุชุญู ูู ุงููู ุงุฐุฌ | |
loaded_models = {} | |
for name, model_id in models.items(): | |
tokenizer = AutoTokenizer.from_pretrained(model_id) | |
model = AutoModelForCausalLM.from_pretrained( | |
model_id, | |
device_map="auto", | |
torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32 | |
) | |
loaded_models[name] = (tokenizer, model) | |
# ุฏุงูุฉ ุงูุชูููุฏ | |
def generate_code(prompt, model_name): | |
tokenizer, model = loaded_models[model_name] | |
inputs = tokenizer(prompt, return_tensors="pt").to(model.device) | |
outputs = model.generate(**inputs, max_new_tokens=150) | |
return tokenizer.decode(outputs[0], skip_special_tokens=True) | |
# ูุงุฌูุฉ Gradio | |
demo = gr.Interface( | |
fn=generate_code, | |
inputs=[ | |
gr.Textbox(lines=5, label="ุงูุชุจ ูุตู ุงูููุฏ (ุจุงูุฅูุฌููุฒูุฉ)"), | |
gr.Radio(choices=list(models.keys()), label="ุงุฎุชุฑ ุงููู ูุฐุฌ") | |
], | |
outputs=gr.Code(label="ุงูููุฏ ุงููุงุชุฌ"), | |
title="Code Generation with Open AI Models", | |
description="ุงุฎุชุฑ ูู ูุฐุฌูุง ู ูุชูุญูุง ูุฃุฏุฎู ูุตููุง ููุชู ุชูููุฏ ุงูููุฏ ุชููุงุฆููุง" | |
) | |
demo.launch() | |