Update app.py

app.py

@@ -1,71 +1,142 @@
-# app.py
-
-from transformers import pipeline
-
-# Load Zero-Shot Classification Model (for detecting Requirement Type, Domain, Stakeholders, and Defects)
-classifier = pipeline("zero-shot-classification", model="facebook/bart-large-mnli")
-
-# Load T5 model for Rewriting (Paraphrasing)
-from transformers import T5ForConditionalGeneration, T5Tokenizer
-
-t5_model = T5ForConditionalGeneration.from_pretrained("t5-small")
-t5_tokenizer = T5Tokenizer.from_pretrained("t5-small")
-
-# Function to classify requirement type (Functional/Non-Functional)
-def classify_requirement_type(requirement):
-    candidate_labels = ["Functional", "Non-Functional"]
-    result = classifier(requirement, candidate_labels)
-    return result['labels'][0]
-
-# Function to identify stakeholders
-def identify_stakeholders(requirement):
-    candidate_labels = ["End User", "Developer", "System Analyst", "Project Manager"]
-    result = classifier(requirement, candidate_labels)
-    return result['labels'][0]
-
-# Function to classify domain of the requirement
-def classify_domain(requirement):
-    candidate_labels = ["Bank", "Healthcare", "Education", "Finance", "Cybersecurity", "E-commerce"]
-    result = classifier(requirement, candidate_labels)
-    return result['labels'][0]
-
-# Function to detect defects (e.g., Ambiguity, Incompleteness)
-def detect_defects(requirement):
-    candidate_labels = ["Ambiguity", "Incompleteness", "Security Flaw", "Redundancy", "Performance Issue"]
-    result = classifier(requirement, candidate_labels)
-    return result['labels'][0]
-
-# Function to rewrite the requirement in a simpler way using T5
-def rewrite_requirement(requirement):
-    input_text = "paraphrase: " + requirement
-    input_ids = t5_tokenizer.encode(input_text, return_tensors="pt", max_length=512, truncation=True)
-    
-    output_ids = t5_model.generate(input_ids, max_length=150, num_beams=5, early_stopping=True)
-    output_text = t5_tokenizer.decode(output_ids[0], skip_special_tokens=True)
-    
-    return output_text
-
-# Main function to take input requirement and process it
-def process_requirement(requirement):
-    requirement_type = classify_requirement_type(requirement)
-    stakeholder = identify_stakeholders(requirement)
-    domain = classify_domain(requirement)
-    defects = detect_defects(requirement)
-    rewritten_requirement = rewrite_requirement(requirement)
-    
+import gradio as gr
+import requests
+from fpdf import FPDF
+import pdfplumber
+
+# Mistral API key (replace with your key)
+API_KEY = "gz6lDXokxgR6cLY72oomALWcm7vhjzQ"
+MISTRAL_API_URL = "https://api.mistral.ai/v1/chat/completions"
+
+# Function to call Mistral API
+def call_mistral_api(prompt):
+    headers = {
+        "Authorization": f"Bearer {API_KEY}",
+        "Content-Type": "application/json"
+    }
+    payload = {
+        "model": "mistral-medium",
+        "messages": [
+            {"role": "user", "content": prompt}
+        ]
+    }
+    response = requests.post(MISTRAL_API_URL, headers=headers, json=payload)
+    if response.status_code == 200:
+        return response.json()['choices'][0]['message']['content']
+    else:
+        return f"Error: {response.status_code}, {response.text}"
+
+# Function to analyze a single requirement
+def analyze_requirement(requirement):
+    # Detect requirement type
+    type_prompt = f"Classify the following requirement as Functional or Non-Functional:\n\n{requirement}\n\nType:"
+    req_type = call_mistral_api(type_prompt)
+
+    # Identify stakeholders
+    stakeholders_prompt = f"Identify the stakeholders for the following requirement:\n\n{requirement}\n\nStakeholders:"
+    stakeholders = call_mistral_api(stakeholders_prompt)
+
+    # Classify domain
+    domain_prompt = f"Classify the domain for the following requirement (e.g., Bank, Healthcare, etc.):\n\n{requirement}\n\nDomain:"
+    domain = call_mistral_api(domain_prompt)
+
+    # Detect defects
+    defects_prompt = f"Detect any defects in the following requirement (e.g., Ambiguity, Incompleteness, etc.):\n\n{requirement}\n\nDefects:"
+    defects = call_mistral_api(defects_prompt)
+
+    # Rewrite requirement
+    rewrite_prompt = f"Rewrite the following requirement in a simpler form:\n\n{requirement}\n\nSimplified Requirement:"
+    rewritten = call_mistral_api(rewrite_prompt)
+
     return {
-        "Requirement Type": requirement_type,
-        "Stakeholder": stakeholder,
+        "Requirement": requirement,
+        "Type": req_type,
+        "Stakeholders": stakeholders,
         "Domain": domain,
         "Defects": defects,
-        "Rewritten Requirement": rewritten_requirement
+        "Rewritten": rewritten
     }
 
-if __name__ == "__main__":
-    # Example usage:
-    requirement = input("Enter the software requirement: ")
-    result = process_requirement(requirement)
-    
-    # Print the results
-    for key, value in result.items():
-        print(f"{key}: {value}")
+# Function to process multiple requirements
+def process_requirements(input_text, uploaded_file):
+    requirements = []
+
+    # If a file is uploaded, extract text from the PDF
+    if uploaded_file:
+        with pdfplumber.open(uploaded_file) as pdf:
+            for page in pdf.pages:
+                requirements.extend(page.extract_text().split("\n"))
+
+    # If text input is provided, split into individual requirements
+    if input_text:
+        requirements.extend(input_text.split("\n"))
+
+    # Analyze each requirement
+    results = []
+    for req in requirements:
+        if req.strip():  # Ignore empty lines
+            results.append(analyze_requirement(req.strip()))
+
+    return results
+
+# Function to generate a PDF report
+def generate_pdf_report(results):
+    pdf = FPDF()
+    pdf.add_page()
+    pdf.set_font("Arial", size=12)
+
+    for result in results:
+        pdf.cell(200, 10, txt=f"Requirement: {result['Requirement']}", ln=True)
+        pdf.cell(200, 10, txt=f"Type: {result['Type']}", ln=True)
+        pdf.cell(200, 10, txt=f"Stakeholders: {result['Stakeholders']}", ln=True)
+        pdf.cell(200, 10, txt=f"Domain: {result['Domain']}", ln=True)
+        pdf.cell(200, 10, txt=f"Defects: {result['Defects']}", ln=True)
+        pdf.cell(200, 10, txt=f"Rewritten: {result['Rewritten']}", ln=True)
+        pdf.cell(200, 10, txt="-" * 50, ln=True)
+
+    pdf_output = "requirements_report.pdf"
+    pdf.output(pdf_output)
+    return pdf_output
+
+# Gradio interface
+def app_interface(input_text, uploaded_file):
+    results = process_requirements(input_text, uploaded_file)
+    pdf_report = generate_pdf_report(results)
+
+    # Display results in a clean format
+    output_text = ""
+    for result in results:
+        output_text += f"Requirement: {result['Requirement']}\n"
+        output_text += f"Type: {result['Type']}\n"
+        output_text += f"Stakeholders: {result['Stakeholders']}\n"
+        output_text += f"Domain: {result['Domain']}\n"
+        output_text += f"Defects: {result['Defects']}\n"
+        output_text += f"Rewritten: {result['Rewritten']}\n"
+        output_text += "-" * 50 + "\n\n"
+
+    return output_text, pdf_report
+
+# Gradio app
+title = "Requirement Analysis App"
+description = """
+This app analyzes software requirements using the Mistral model. 
+You can either enter requirements manually or upload a PDF file containing requirements.
+"""
+team_name = "Team Innovators"
+
+interface = gr.Interface(
+    fn=app_interface,
+    inputs=[
+        gr.Textbox(label="Enter Requirements (one per line)"),
+        gr.File(label="Upload PDF File (optional)")
+    ],
+    outputs=[
+        gr.Textbox(label="Analysis Results"),
+        gr.File(label="Download PDF Report")
+    ],
+    title=title,
+    description=description,
+    theme="soft"
+)
+
+# Add team name to the interface
+interface.launch()