fix token problem

src/compliance_lib.py

@@ -1,40 +1,37 @@
 import re
-from huggingface_hub import InferenceClient
 import os
 import requests as req
 from bs4 import BeautifulSoup
-import streamlit as st
-from dotenv import load_dotenv
-load_dotenv()
-RULES={
-    "GDPR":[
-     ("Lawful basis documented", r"lawful\s+basis"),
+from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
+import torch
+
+RULES = {
+    "GDPR": [
+        ("Lawful basis documented", r"lawful\s+basis"),
         ("Data-subject rights process", r"right\s+to\s+access|erasure"),
         ("72-hour breach notice plan", r"72\s*hour"),
     ],
-    "EU_AI_ACT":[
-          ("High-risk AI DPIA", r"risk\s+assessment"),
+    "EU_AI_ACT": [
+        ("High-risk AI DPIA", r"risk\s+assessment"),
         ("Training data governance", r"data\s+governance"),
     ],
-    "ISO_27001":[
+    "ISO_27001": [
         ("Annex A control list", r"annex\s*a"),
         ("Statement of Applicability", r"statement\s+of\s+applicability"),
     ]
-    }
+}
+
 
-def run_check(text,framework):
-    # print(text,framework) #array from me aata hai framework
-    results={}
+def run_check(text, framework):
+    results = {}
     for fw in framework:
-        results[fw]=[]  #store particular fw data
-            # one work as label & one work as pattern e.g==>label: Training data governance pattern: data\s+governance
+        results[fw] = []
         for label, pattern in RULES[fw]:
-           match = re.search(pattern, text, re.I)  # re.I = re.IGNORECASE
-           results[fw].append((label, bool(match)))
+            match = re.search(pattern, text, re.I)
+            results[fw].append((label, bool(match)))
     return results
 
 
-
 AI_REPORT_PROMPT = """
 You are an expert compliance consultant with deep experience in GDPR, the EU AI Act, ISO 27001, and related global data‑privacy and security standards. You have just received a concise checklist summary showing, for each framework, how many controls passed and which specific items failed.
 
@@ -75,32 +72,39 @@ Generate the report as markdown.
 
 HF_MODEL = "mistralai/Mixtral-8x7B-Instruct-v0.1"
 
-
-
-
-def generate_report(prompt,max_tokens=600):
-     token = os.getenv("HF_TOKEN")
-     if not token:
-        raise EnvironmentError("token is not found in env issue")
-    
-     client = InferenceClient(
-    provider="together",
-    api_key=token,
-)    
-     try:
-      response = client.chat.completions.create(
-            model=HF_MODEL,
-            messages=[ {
-            "role": "user",
-            "content": prompt
-        }]
-          
+# Load the text generation pipeline
+def load_pipeline():
+    tokenizer = AutoTokenizer.from_pretrained(HF_MODEL)
+    model = AutoModelForCausalLM.from_pretrained(
+        HF_MODEL,
+        torch_dtype=torch.float16,
+        device_map="auto",
+        trust_remote_code=True
+    )
+    pipe = pipeline(
+        "text-generation",
+        model=model,
+        tokenizer=tokenizer,
+        device_map="auto"
+    )
+    return pipe
+
+generator = load_pipeline()
+
+
+def generate_report(prompt, max_tokens=600):
+    try:
+        response = generator(
+            prompt,
+            max_new_tokens=max_tokens,
+            do_sample=True,
+            temperature=0.7,
+            top_p=0.95,
+            return_full_text=False
         )
-    
-      return response.choices[0].message.content
-     except Exception as e:
-        
-        return "Error: Failed to generate report."
+        return response[0]["generated_text"]
+    except Exception as e:
+        return f"Error: {e}"
 
 
 def fetchText(url):
@@ -113,8 +117,9 @@ def fetchText(url):
             text = main_content.get_text(separator='\n', strip=True)
         else:
             text = soup.body.get_text(separator='\n', strip=True)
-        
-        return text.strip(), None  # No error
+        return text.strip(), None
     except Exception as e:
         return "", f"Error fetching URL: {e}"
-__all__=["RULES","run_check","AI_REPORT_PROMPT","generate_report","fetchText"]
+
+# Exported functions
+__all__ = ["RULES", "run_check", "AI_REPORT_PROMPT", "generate_report", "fetchText"]