Create app.py

app.py

@@ -0,0 +1,88 @@
+import gradio as gr
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from peft import PeftModel
+
+# --- Configuration ---
+# ⚠️ FIXED USERNAME
+BASE_MODEL_ID = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
+ADAPTER_MODEL_ID = "Vivek16/Root_Math-TinyLlama-CPU" 
+
+# Define the instruction template used during fine-tuning (Step 5)
+INSTRUCTION_TEMPLATE = "<|system|>\nSolve the following math problem:</s>\n<|user|>\n{}</s>\n<|assistant|>"
+
+# --- Model Loading Function ---
+def load_model():
+    """Loads the base model and merges the LoRA adapters."""
+    print("Loading base model...")
+    tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL_ID)
+    model = AutoModelForCausalLM.from_pretrained(
+        BASE_MODEL_ID, 
+        torch_dtype=torch.bfloat16, # Use bfloat16 for efficiency on CPU
+        device_map="cpu"
+    )
+
+    print("Loading and merging PEFT adapters...")
+    # Load the trained LoRA adapters from your repo
+    model = PeftModel.from_pretrained(model, ADAPTER_MODEL_ID)
+    # Merge the adapter weights into the base model weights
+    model = model.merge_and_unload()
+    model.eval()
+
+    # Ensure pad token is set for generation
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+
+    print("Model loaded and merged successfully!")
+    return tokenizer, model
+
+# Load the model outside the prediction function for efficiency
+tokenizer, model = load_model()
+
+# --- Prediction Function ---
+def generate_response(prompt):
+    """Generates a response using the fine-tuned model."""
+    # 1. Format the user input using the exact chat template
+    formatted_prompt = INSTRUCTION_TEMPLATE.format(prompt)
+
+    # 2. Tokenize the input
+    inputs = tokenizer(formatted_prompt, return_tensors="pt")
+
+    # 3. Generate the response (on CPU)
+    with torch.no_grad():
+        output_tokens = model.generate(
+            **inputs,
+            max_new_tokens=256,
+            do_sample=True,
+            temperature=0.7,
+            top_k=50,
+            pad_token_id=tokenizer.eos_token_id 
+        )
+
+    # 4. Decode the output and strip the prompt
+    generated_text = tokenizer.decode(output_tokens[0], skip_special_tokens=False)
+
+    # Extract only the assistant's response (everything after the last <|assistant|> tag)
+    response_start = generated_text.rfind('<|assistant|>')
+    if response_start != -1:
+        # Get the text after <|assistant|> and strip the trailing </s>
+        assistant_response = generated_text[response_start + len('<|assistant|>'):].strip().split('</s>')[0].strip()
+    else:
+        assistant_response = "Error: Could not parse model output."
+
+    return assistant_response
+
+# --- Gradio Interface ---
+title = "Root Math TinyLlama 1.1B - CPU Fine-Tuned"
+description = "A CPU-friendly TinyLlama model fine-tuned on the Big-Math-RL-Verified dataset using LoRA."
+article = "Model repository: Vivek16/Root_Math-TinyLlama-CPU"
+
+gr.Interface(
+    fn=generate_response,
+    inputs=gr.Textbox(lines=5, label="Enter your math problem here:"),
+    outputs=gr.Textbox(label="Model Answer"),
+    title=title,
+    description=description,
+    article=article,
+    theme="soft"
+).launch()