athul020's picture
Upload folder using huggingface_hub
5a043a6 verified
Raw
History Blame Contribute Delete
1.94 kB
import gradio as gr
from infinite_context import ContextGateway
# Initialize the Context Gateway with a lightweight model by default
# On a Space with a GPU, it will automatically use Unsloth optimization.
gateway = ContextGateway(model_id="Qwen/Qwen2.5-0.5B-Instruct")
def ask_question(context: str, question: str) -> str:
if not context.strip():
return "Please provide a context document."
if not question.strip():
return "Please ask a question."
try:
# Memorize the context using Test-Time Training (in-place)
gateway.memorise(context, in_place=True)
# Ask the question based on the memorized context
response = gateway.ask(question)
return response
except Exception as e:
return f"An error occurred: {str(e)}"
with gr.Blocks(title="Infinite Context - TTT Memory") as app:
gr.Markdown("# Infinite Context 🧠")
gr.Markdown("Inject unlimited context into LLMs using Test-Time Training (TTT) Fast Weights. Paste your large document below, and the model will learn it dynamically.")
with gr.Row():
with gr.Column(scale=2):
context_box = gr.Textbox(
label="Context Document",
lines=15,
placeholder="Paste your long document, codebase, or book here...",
)
with gr.Column(scale=1):
question_box = gr.Textbox(
label="Question",
lines=3,
placeholder="What is this document about?",
)
submit_btn = gr.Button("Ask", variant="primary")
output_box = gr.Textbox(
label="Answer",
lines=8,
interactive=False,
)
submit_btn.click(
fn=ask_question,
inputs=[context_box, question_box],
outputs=[output_box],
)
if __name__ == "__main__":
app.launch()