app.py

import gradio as gr
import time

from transformers import AutoTokenizer, AutoModelForCausalLM, TextGenerationPipeline

MODEL_ID = "akshaynayaks9845/rml-ai-phi1_5-rml-100k"

def load_pipeline():
    try:
        tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
        model = AutoModelForCausalLM.from_pretrained(MODEL_ID)
        pipe = TextGenerationPipeline(model=model, tokenizer=tokenizer, device=-1)
        return pipe
    except Exception as e:
        return str(e)

pipe_or_err = load_pipeline()

SAMPLES = [
    "What is artificial intelligence?",
    "Explain machine learning in one sentence.",
    "What is quantum computing?",
]

def generate_response(prompt, max_new_tokens=128, temperature=0.2):
    start = time.time()
    if isinstance(pipe_or_err, str):
        return f"Model load error: {pipe_or_err}"
    try:
        outputs = pipe_or_err(
            prompt,
            max_new_tokens=int(max_new_tokens),
            do_sample=bool(temperature and temperature > 0),
            temperature=float(temperature),
            top_p=0.9,
            repetition_penalty=1.1,
            truncation=True,
        )
        text = outputs[0]["generated_text"]
        # Return only continuation if the model echoes the prompt
        reply = text[len(prompt):].strip() if text.startswith(prompt) else text
        elapsed = int((time.time() - start) * 1000)
        return f"{reply}

(⏱️ {elapsed} ms)"
    except Exception as e:
        return f"Error: {str(e)}"

with gr.Blocks(title="RML-AI Demo") as demo:
    gr.Markdown('''
    # RML-AI Demo
    Ask a question below. The model will respond in GPT-style. This is a lightweight prototype demo.
    ''')
    with gr.Row():
        prompt = gr.Textbox(label="Your question", value=SAMPLES[0])
    with gr.Row():
        max_new = gr.Slider(32, 256, value=128, step=16, label="Max new tokens")
        temp = gr.Slider(0.0, 1.0, value=0.2, step=0.1, label="Temperature")
    with gr.Row():
        btn = gr.Button("Generate")
    output = gr.Textbox(label="Answer", lines=8)
    with gr.Row():
        gr.Examples(SAMPLES, inputs=prompt)

    btn.click(generate_response, [prompt, max_new, temp], output)

if __name__ == "__main__":
    demo.launch()