ModelSpace / app.py
AxisCommunity's picture
Update app.py
ffbb716 verified
raw
history blame contribute delete
869 Bytes
import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
import torch
# Загружаем модель с параметрами для слабых систем
model_name = "AxisCommunity/OrionPaxAI_1.0V"
# Используем torch_dtype=torch.float16, чтобы уменьшить вес в памяти в 2 раза
# device_map="auto" поможет распределить нагрузку
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
model_name,
torch_dtype=torch.float16,
device_map="auto",
low_cpu_mem_usage=True
)
pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
def predict(text):
return pipe(text, max_new_tokens=50)[0]['generated_text']
demo = gr.Interface(fn=predict, inputs="text", outputs="text")
demo.launch()