rubpy-coder / app.py
Mirdehghan's picture
Update app.py
320d2f8 verified
raw
history blame contribute delete
654 Bytes
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
import gradio as gr
MODEL_PATH = "./rubpy-model"
tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
model = AutoModelForCausalLM.from_pretrained(
MODEL_PATH,
device_map="auto"
)
def chat(prompt):
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
output = model.generate(
**inputs,
max_new_tokens=512,
temperature=0.3
)
return tokenizer.decode(output[0], skip_special_tokens=True)
gr.Interface(
fn=chat,
inputs=gr.Textbox(lines=5),
outputs=gr.Textbox(lines=15),
title="Rubpy AI Coder"
).launch()