GenueApp / app.py
Scoopala7's picture
fix model handling + more part 2
c2aaf98 verified
Raw
History Blame
2.2 kB
import os
import gradio as gr
from transformers import pipeline
MODEL_ID = os.getenv("MODEL_ID", "GenueAI/Inelly-4.5-Blaze")
MODEL_NAME = os.getenv("MODEL_NAME", "Inelly 4.5 Blaze")
pipe = None
def load_model():
global pipe
pipe = pipeline(
"text-generation",
model=MODEL_ID,
torch_dtype="auto",
model_kwargs={
"low_cpu_mem_usage": True,
"device_map": "sequential"
}
)
def build_prompt(message_text, history, system_prompt):
messages = []
if system_prompt.strip():
messages.append({"role": "system", "content": system_prompt.strip()})
for msg in history:
messages.append({"role": msg["role"], "content": msg["content"]})
messages.append({"role": "user", "content": message_text})
return messages
def chat(message, history, system_prompt):
message_text = message.get("text", "").strip() if isinstance(message, dict) else str(message).strip()
if not message_text:
return ""
prompt = build_prompt(message_text, history, system_prompt)
outputs = pipe(prompt)
try:
return outputs[0]["generated_text"][-1]["content"]
except (KeyError, IndexError, TypeError):
try:
return outputs["generated_text"][-1]["content"]
except (KeyError, IndexError, TypeError):
return str(outputs)
with gr.Blocks(title="Genue Chat") as demo:
gr.Markdown("# Genue Chat")
gr.Markdown(f"Chat with {MODEL_NAME} from Hugging Face.")
with gr.Row():
with gr.Column(scale=4):
chatbot = gr.ChatInterface(
fn=chat,
additional_inputs=[
gr.Textbox(
label="System prompt",
value="You are a helpful assistant.",
lines=3,
),
],
textbox=gr.Textbox(
label="Prompt",
placeholder="Ask anything...",
container=False,
scale=7,
),
)
if __name__ == "__main__":
load_model()
demo.queue()
demo.launch()