Update app.py

app.py

@@ -1,62 +1,62 @@
-from huggingface_hub import hf_hub_download
-from llama_cpp import Llama
+import json
 import gradio as gr
+from llama_cpp import Llama
+from huggingface_hub import hf_hub_download
+import os
+import re
+
+# Download your GGUF file from your Hugging Face repo
+MODEL_NAME = "llama3-pr-Q4_K_M.gguf"
+MODEL_REPO = "Bton/llama3-pr-Q4_K_M.gguf"
 
-# Download the GGUF model from Hugging Face Hub
-model_path = hf_hub_download(
-    repo_id="Bton/llama3-product-reviewer",  # ✅ Make sure this matches your repo name
-    filename="unsloth.Q4_K_M.gguf",
+local_model_path = hf_hub_download(
+    repo_id=MODEL_REPO,
+    filename=MODEL_NAME,
     local_dir="."
 )
 
-# Initialize the LLM
+# Load model with llama-cpp
 llm = Llama(
-    model_path=model_path,
+    model_path=local_model_path,
+    n_gpu_layers=-1,  # GPU offload if available
     n_ctx=2048,
-    n_threads=4,
-    n_gpu_layers=35,
-    use_mlock=True,
-    verbose=False
 )
 
-# Prompt Template
-TEMPLATE = """<|im_start|>system
-You are a helpful product review generator. Given Amazon product details, write a convincing, natural-sounding review in JSON format with "title" and "review" fields.
-<|im_end|>
-<|im_start|>user
-Write a product review based on the following product details.
-Your input is:
-{input_json}
-<|im_end|>
-<|im_start|>assistant
-"""
+# Safe JSON extractor
+def safe_parse_json(text):
+    try:
+        return json.loads(text.strip())
+    except json.JSONDecodeError:
+        try:
+            match = re.search(r'\{.*\}', text, re.DOTALL)
+            if match:
+                return json.loads(match.group())
+        except:
+            pass
+    return {
+        "title": "Generated Title",
+        "review": "Error parsing response"
+    }
 
 # Inference function
-def generate_review(input_json):
-    full_prompt = TEMPLATE.format(input_json=input_json.strip())
-
-    output = llm(
-        full_prompt,
-        max_tokens=512,
-        temperature=0.7,
-        stop=["<|im_end|>", "<|end_of_text|>"]
-    )
-
-    try:
-        response = output["choices"][0]["text"].strip()
-        parsed = eval(response)  # only safe here since you control the model format
-        return parsed.get("title", ""), parsed.get("review", "")
-    except Exception as e:
-        return "Error parsing response", str(e)
+def generate_review(input_data):
+    prompt = f"{input_data}"
+    response = llm(prompt, max_tokens=512, stop=["<|eot_id|>"], stream=False)
+    full_output = response["choices"][0]["text"]
+    parsed = safe_parse_json(full_output)
+    return parsed["title"], parsed["review"]
 
 # Gradio UI
-with gr.Blocks() as demo:
-    gr.Markdown("## 🛍️ Amazon Product Review Generator (GGUF on Llama.cpp)")
-    input_json = gr.Textbox(label="Paste JSON product details here", lines=8, placeholder='{"product_title": "...", "price": "...", ... }')
-    btn = gr.Button("Generate Review")
-    title_out = gr.Textbox(label="Generated Title")
-    review_out = gr.Textbox(label="Generated Review", lines=5)
-
-    btn.click(generate_review, inputs=input_json, outputs=[title_out, review_out])
+demo = gr.Interface(
+    fn=generate_review,
+    inputs=gr.Textbox(label="Product Info (JSON format)"),
+    outputs=[
+        gr.Textbox(label="Generated Title"),
+        gr.Textbox(label="Generated Review", lines=6)
+    ],
+    title="🛒 Product Review Generator",
+    description="Paste product info as JSON and generate a helpful review using your fine-tuned GGUF model.",
+)
 
-demo.launch()
+if __name__ == "__main__":
+    demo.launch()