Update app.py

app.py

@@ -1,96 +1,80 @@
-# app.py - PromptPolice (Local Transformers Version)
+# 👮 PromptPolice MVP - Streamlit App with Mistral Backend (JSON Output)
+
 import streamlit as st
-import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
+import requests
 import json
-from huggingface_hub import login
-import os
 
-login(os.environ["HF_PROJECT_TOKEN"])
+HF_API_URL = "https://router.huggingface.co/novita/v3/openai/chat/completions"
+HF_TOKEN = "HF_PROJECT_TOKEN"
 
-st.set_page_config(page_title="👮 PromptPolice", layout="centered")
-st.title("👮 PromptPolice: Prompt Evaluator")
+HEADERS = {
+    "Authorization": f"Bearer {HF_TOKEN}",
+    "Content-Type": "application/json"
+}
+
+PROMPT_TEMPLATE = """
+You are a prompt evaluation assistant. Evaluate the following user prompt in JSON format using the structure provided below.
 
-PROMPT_TEMPLATE = """You are a prompt evaluation assistant. Evaluate the following user prompt in JSON format using the structure provided below.
 Prompt:
+
 {user_prompt}
+
 Evaluate based on the following criteria:
 - Clarity (1-5)
 - Context (1-5)
 - Specificity (1-5)
 - Intent (1-5)
 Also include a suggestion for improving the prompt.
+
 Respond ONLY in this JSON format:
-{{
+{
   "prompt": "...",
-  "evaluation": {{
+  "evaluation": {
     "Clarity": ..., 
     "Context": ..., 
     "Specificity": ..., 
     "Intent": ..., 
     "suggestion": "..."
-  }}
-}}"""
+  }
+}
+"""
 
+def evaluate_prompt(user_prompt):
+    payload = {
+        "inputs": PROMPT_TEMPLATE.format(user_prompt=user_prompt),
+        "parameters": {"max_new_tokens": 512, "temperature": 0.7}
+    }
 
+    response = requests.post(HF_API_URL, headers=HEADERS, json=payload)
+    result = response.json()
 
-# --- Load Model and Tokenizer ---
-@st.cache_resource
-def load_model():
-    model_id = "mistralai/Mistral-7B-Instruct-v0.1"
-
-    bnb_config = BitsAndBytesConfig(
-        load_in_4bit=True,
-        bnb_4bit_compute_dtype=torch.float16,
-        bnb_4bit_use_double_quant=True,
-        bnb_4bit_quant_type="nf4"
-    )
-
-    tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
-    model = AutoModelForCausalLM.from_pretrained(
-        model_id,
-        device_map="cpu",
-        torch_dtype=torch.float32  # Use float32 if float16 fails
-    )
-    return tokenizer, model
+    # Handle streaming/text output
+    if isinstance(result, list) and "generated_text" in result[0]:
+        raw_text = result[0]["generated_text"]
+        try:
+            # Try to parse the JSON segment only
+            json_part = raw_text[raw_text.index("{"):]
+            return json.loads(json_part)
+        except:
+            return {"error": "Failed to parse model output."}
+    else:
+        return result
 
-tokenizer, model = load_model()
+# --- Streamlit UI ---
+st.set_page_config(page_title="👮 PromptPolice", layout="centered")
+st.title("👮 PromptPolice: Prompt Evaluator")
 
-# --- Prompt Input ---
 user_prompt = st.text_area("Paste your AI prompt here:", height=200)
 
-def evaluate_prompt_local(prompt_text):
-    full_prompt = PROMPT_TEMPLATE.format(user_prompt=prompt_text)
-    inputs = tokenizer(full_prompt, return_tensors="pt").to(model.device)
-    outputs = model.generate(
-        **inputs,
-        max_new_tokens=512,
-        temperature=0.7,
-        repetition_penalty=1.1,
-        eos_token_id=tokenizer.eos_token_id,
-    )
-    decoded = tokenizer.decode(outputs[0][inputs["input_ids"].shape[-1]:], skip_special_tokens=True)
-
-    # Try extracting only the JSON response from the output
-    try:
-        json_part = decoded[decoded.index("{"):decoded.rindex("}")+1]
-        return json.loads(json_part)
-    except Exception as e:
-        return {"error": f"Failed to parse model output: {str(e)}", "raw": decoded}
-
-# --- Run Evaluation ---
 if st.button("Evaluate Prompt") and user_prompt:
-    with st.spinner("Evaluating your prompt locally with Mistral..."):
-        evaluation_result = evaluate_prompt_local(user_prompt)
+    with st.spinner("Evaluating your prompt with Mistral..."):
+        evaluation_result = evaluate_prompt(user_prompt)
 
-    if "error" in evaluation_result:
-        st.error(evaluation_result["error"])
-        st.text_area("Raw Output:", value=evaluation_result.get("raw", ""), height=250)
-    else:
-        st.subheader("Evaluation Result (JSON):")
-        st.code(json.dumps(evaluation_result, indent=2), language='json')
+    st.subheader("Evaluation Result (JSON):")
+    st.code(json.dumps(evaluation_result, indent=2), language='json')
 
-        if st.button("💾 Save to Dataset"):
-            with open("crowdsourced_prompts.jsonl", "a", encoding="utf-8") as f:
-                f.write(json.dumps(evaluation_result) + "\n")
-            st.success("Prompt evaluation saved!")
+    # Optional: Append to local file for crowd-sourced fine-tuning later
+    if st.button("💾 Save to Dataset"):
+        with open("crowdsourced_prompts.jsonl", "a", encoding="utf-8") as f:
+            f.write(json.dumps(evaluation_result) + "\n")
+        st.success("Prompt evaluation saved!")