Spaces:
Build error
Build error
| 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() | |