Update app.py

app.py

@@ -1,70 +1,88 @@
 import gradio as gr
-from huggingface_hub import InferenceClient
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from peft import PeftModel
 
+# --- Configuration ---
+# ⚠️ REPLACE 'YOUR_HF_USERNAME' with your actual username
+BASE_MODEL_ID = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
+ADAPTER_MODEL_ID = "YOUR_HF_USERNAME/Root_Math-TinyLlama-CPU" 
 
-def respond(
-    message,
-    history: list[dict[str, str]],
-    system_message,
-    max_tokens,
-    temperature,
-    top_p,
-    hf_token: gr.OAuthToken,
-):
-    """
-    For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
-    """
-    client = InferenceClient(token=hf_token.token, model="openai/gpt-oss-20b")
+# 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|>"
 
-    messages = [{"role": "system", "content": system_message}]
+# --- 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"
+    )
 
-    messages.extend(history)
+    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
 
-    messages.append({"role": "user", "content": message})
+# Load the model outside the prediction function for efficiency
+tokenizer, model = load_model()
 
-    response = ""
+# --- 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")
 
-    for message in client.chat_completion(
-        messages,
-        max_tokens=max_tokens,
-        stream=True,
-        temperature=temperature,
-        top_p=top_p,
-    ):
-        choices = message.choices
-        token = ""
-        if len(choices) and choices[0].delta.content:
-            token = choices[0].delta.content
+    # 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
 
-        response += token
-        yield 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: " + ADAPTER_MODEL_ID
 
-
-"""
-For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
-"""
-chatbot = gr.ChatInterface(
-    respond,
-    type="messages",
-    additional_inputs=[
-        gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
-        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
-        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
-        gr.Slider(
-            minimum=0.1,
-            maximum=1.0,
-            value=0.95,
-            step=0.05,
-            label="Top-p (nucleus sampling)",
-        ),
-    ],
-)
-
-with gr.Blocks() as demo:
-    with gr.Sidebar():
-        gr.LoginButton()
-    chatbot.render()
-
-
-if __name__ == "__main__":
-    demo.launch()
+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()