gemma_api / app.py
gg34455's picture
Update app.py
d3419fc verified
Raw
History Blame Contribute Delete
747 Bytes
import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
MODEL_ID = "Jackrong/Qwen3.5-9B-GLM5.1-Distill-v1"
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
model = AutoModelForCausalLM.from_pretrained(
MODEL_ID,
device_map="cpu",
torch_dtype=torch.float32,
low_cpu_mem_usage=True
)
def chat(prompt):
inputs = tokenizer(prompt, return_tensors="pt")
outputs = model.generate(
**inputs,
max_new_tokens=64, # keep LOW
do_sample=True,
temperature=0.7
)
return tokenizer.decode(outputs[0], skip_special_tokens=True)
app = gr.Interface(
fn=chat,
inputs="text",
outputs="text",
api_name="generate"
)
app.queue()
app.launch()