# app.py import os import gradio as gr from huggingface_hub import InferenceApi from typing import List, Tuple HF_TOKEN = os.environ.get("HUGGINGFACEHUB_API_TOKEN") if not HF_TOKEN: raise RuntimeError("Set HUGGINGFACEHUB_API_TOKEN in your Space secrets/settings.") # Default models (you can replace these with your own hosted model names) TEXT_MODEL = os.environ.get("YAGA_TEXT_MODEL", "mistralai/Mistral-7B-v0.1") # hosted generation model IMAGE_CAPTION_MODEL = os.environ.get("YAGA_IMAGE_MODEL", "Salesforce/blip-image-captioning-large") WHISPER_MODEL = os.environ.get("YAGA_WHISPER_MODEL", "openai/whisper-large-v2") # Inference API clients text_api = InferenceApi(repo_id=TEXT_MODEL, token=HF_TOKEN) image_api = InferenceApi(repo_id=IMAGE_CAPTION_MODEL, token=HF_TOKEN) whisper_api = InferenceApi(repo_id=WHISPER_MODEL, token=HF_TOKEN) SYSTEM_PROMPT = ( "You are YAGA — an advanced, helpful, and safe AI assistant. " "Answer concisely and clearly. If user provided an image or audio, " "use the provided descriptions/transcripts as context." ) # Simple in-memory conversation store per session id (Gradio session) # For Spaces this is ephemeral; for production, persist to DB. CONV_STORE = {} def build_prompt(history: List[Tuple[str, str]], extra_context: str | None): parts = [f"SYSTEM: {SYSTEM_PROMPT}"] if extra_context: parts.append(f"CONTEXT: {extra_context}") for role, text in history: parts.append(f"{role}: {text}") parts.append("Assistant:") return "\n".join(parts) def get_session_history(session_id: str): return CONV_STORE.setdefault(session_id, []) def clear_chat(session_id: str): CONV_STORE[session_id] = [] return [] def transcribe_audio(audio_bytes): # Whisper inference: input is bytes or file-like; InferenceApi accepts bytes try: out = whisper_api(inputs=audio_bytes) # InferenceApi may return plain text or dict depending on model; handle common shapes if isinstance(out, dict) and "text" in out: return out["text"] if isinstance(out, str): return out # fallback: try join if list return str(out) except Exception as e: return f"[transcription_error: {e}]" def caption_image(image_bytes): try: out = image_api(inputs=image_bytes) # Many image caption models return plain text if isinstance(out, dict) and "generated_text" in out: return out["generated_text"] if isinstance(out, str): return out # fallback return str(out) except Exception as e: return f"[caption_error: {e}]" def generate_reply(prompt: str, max_tokens: int = 256, temperature: float = 0.7): params = {"max_new_tokens": max_tokens, "temperature": temperature, "top_p":0.95} # Some Inference API models return dicts or plain text; just return textual result out = text_api(inputs=prompt, parameters=params) if isinstance(out, dict): # common place: {"generated_text": "..."} if "generated_text" in out: return out["generated_text"] # some endpoints return list of candidates if "error" in out: return f"[generation_error: {out['error']}]" try: return str(out) except: return "[generation_error: unexpected format]" if isinstance(out, str): return out return str(out) def respond(session_id, user_message, image_file, audio_file, max_tokens=256, temperature=0.7): history = get_session_history(session_id) extra_ctx_parts = [] if image_file is not None: img_bytes = image_file.read() caption = caption_image(img_bytes) extra_ctx_parts.append(f"Image caption: {caption}") if audio_file is not None: audio_bytes = audio_file.read() transcript = transcribe_audio(audio_bytes) extra_ctx_parts.append(f"Audio transcript: {transcript}") extra_context = "\n".join(extra_ctx_parts) if extra_ctx_parts else None # Append user message to history history.append(("User", user_message)) prompt = build_prompt(history, extra_context) reply = generate_reply(prompt, max_tokens=int(max_tokens), temperature=float(temperature)) history.append(("Assistant", reply)) # Keep last N turns CONV_STORE[session_id] = history[-12*2:] # Convert to chat-like pairs for the frontend chat_pairs = [(r[0], r[1]) for r in CONV_STORE[session_id]] return chat_pairs, "" with gr.Blocks(title="YAGA — HF Space Prototype") as demo: gr.Markdown("## YAGA — Multi-turn chat (Hugging Face Space)\nText, image and audio inputs supported.\n\n**Important:** Add your Hugging Face token to Space secrets as `HUGGINGFACEHUB_API_TOKEN`.") with gr.Row(): with gr.Column(scale=3): chatbot = gr.Chatbot() msg = gr.Textbox(placeholder="Type your message here...", label="Message") with gr.Row(): img = gr.Image(type="bytes", label="Optional image") aud = gr.Audio(source="upload", type="file", label="Optional audio") with gr.Row(): max_t = gr.Slider(32, 1024, value=256, step=32, label="Max new tokens") temp = gr.Slider(0.0, 1.0, value=0.7, step=0.05, label="Temperature") with gr.Row(): send = gr.Button("Send") clear = gr.Button("Clear chat") with gr.Column(scale=1): gr.Markdown("### Session Info") session_id_display = gr.Textbox(value="", interactive=False, label="Session ID (auto)") gr.Markdown("### Hints") gr.Markdown("- Upload an image to give the model visual context (it will be captioned).\n- Upload audio to transcribe via Whisper.\n- Replace `TEXT_MODEL` env var with your model repo for a custom model.") # use gradio session_state for session id def get_sid(): import uuid return str(uuid.uuid4()) sid = get_sid() session_id_display.value = sid def on_send(message, image_file, audio_file, max_tokens, temperature): pairs, _ = respond(sid, message, image_file, audio_file, max_tokens, temperature) # convert pairs to Chatbot messages: list of [user, assistant] pairs chat_display = [] it = iter(pairs) # pairs is list of (role, text) # Build sequential chat list for role, text in pairs: # only show in alternating style: we will append tuple (role, text) chat_display.append((role, text)) # For gradio Chatbot, we need list of [user, assistant] pairs grouped. # But we'll just return the last assistant message to append: last_user = pairs[-2][1] if len(pairs) >= 2 else "" last_assistant = pairs[-1][1] return gr.update(value=chat_display), "" send.click(on_send, [msg, img, aud, max_t, temp], [chatbot, msg]) def on_clear(): clear_chat(sid) return [], "" clear.click(on_clear, [], [chatbot, msg]) if __name__ == "__main__": demo.launch()