app.py

import os
os.environ["CUDA_VISIBLE_DEVICES"] = ""
os.environ['TF_ENABLE_ONEDNN_OPTS'] = '0'
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"

from transformers import RagTokenizer, RagSequenceForGeneration
import gradio as gr
import torch

# Documents from the provided text files
documents = [
    "All employees are expected to adhere to the company’s code of conduct, which emphasizes professionalism, integrity, and respect in the workplace. Arriving on time for work, meetings, and scheduled activities is mandatory. Employees must maintain a professional demeanor when interacting with colleagues, clients, and stakeholders. Respect for diversity and inclusion is a core value, and any form of discrimination or harassment will not be tolerated. Violations of the code of conduct may lead to disciplinary action, including warnings, suspension, or termination, depending on the severity of the offense. Employees are encouraged to report any concerns to their manager or the HR department immediately. The company provides an anonymous reporting hotline for sensitive issues.",
    "Employees are entitled to 20 days of paid annual leave per year to ensure work-life balance. Leave requests must be submitted through the company’s HR portal at least 2 weeks in advance. Managers will review and approve leave based on team workload and business needs. Unused leave can be carried over to the next year, but the maximum carryover is capped at 10 days. Employees are encouraged to plan their leave early to avoid conflicts during peak business periods. Leave beyond the entitled amount may be granted as unpaid leave at the discretion of the manager. Public holidays are observed as per the national calendar, and employees are entitled to these days off with pay. If required to work on a public holiday, employees will receive compensatory time off or overtime pay.",
    "Remote work is permitted for eligible employees with prior manager approval to promote flexibility. Employees must ensure a stable internet connection with a minimum speed of 25 Mbps for video conferencing and collaboration tools. A distraction-free environment is required to maintain productivity during remote work hours. Employees working remotely must adhere to the same working hours as in-office staff, typically 9 AM to 5 PM. All company equipment, such as laptops and monitors, must be used exclusively for work purposes. Remote employees are required to attend virtual team meetings and check-ins as scheduled by their manager. The company may conduct periodic audits to ensure compliance with remote work policies. Any misuse of company resources or failure to meet performance expectations may result in the revocation of remote work privileges.",
    "The company provides a dress code policy to maintain a professional image in the workplace. Business casual attire is required for all employees during office hours. On Fridays, employees may dress in smart casual attire as part of a relaxed dress code day. Employees attending client meetings or external events must wear formal business attire, such as suits or dresses. Visible tattoos and excessive piercings should be covered or minimized during work hours. Footwear must be closed-toe in the office for safety reasons. Non-compliance with the dress code may result in a verbal warning for the first offense, followed by written warnings for repeated violations. Employees are expected to use their best judgment to align with the company’s professional standards.",
    "The company is committed to maintaining a safe and healthy work environment for all employees. Smoking is prohibited in all company premises, including parking areas and break rooms. Employees must follow safety protocols during emergencies, such as fire drills and evacuations. First aid kits are available on every floor, and employees are trained annually on basic first aid procedures. Hazardous materials, if any, must be handled by trained personnel only. Employees are required to report any safety concerns, such as faulty equipment or unsafe conditions, to the facilities team immediately. The company conducts quarterly safety audits to ensure compliance with occupational health and safety regulations. Failure to adhere to safety protocols may result in disciplinary action to protect the well-being of all staff.",
    "All employees are required to sign a non-disclosure agreement (NDA) as a condition of employment with the company. The NDA ensures the protection of confidential information, including trade secrets, client data, and proprietary technology. Employees must not disclose any sensitive information to external parties without prior written consent from the company. This obligation extends to discussions on social media platforms or public forums. Breaches of the NDA will be taken seriously and may result in legal action, including lawsuits for damages. Employees are required to return all confidential materials upon termination of employment. The NDA remains in effect even after an employee leaves the company for a period of 5 years.",
    "The standard employment contract includes a 3-month probation period for all new hires to assess their performance. During this period, employees will receive regular feedback from their managers to support their onboarding. Termination during probation may occur with 2 weeks’ notice if performance does not meet expectations. Employees are entitled to all benefits, including health insurance, during the probation period. Successful completion of the probation period will result in confirmation of permanent employment status. Managers will conduct a formal review meeting at the end of the probation period to discuss performance and set goals. Employees may request an extension of the probation period if additional time is needed to meet performance standards.",
    "Employees are subject to a non-compete clause that prohibits working for competitors for 1 year after termination. This clause applies to direct competitors within the same industry and geographic region as the company. The company defines competitors as any organization offering similar products or services. Employees must disclose any potential conflicts of interest before engaging in external employment or consulting work. Violation of the non-compete clause may result in legal action to enforce the terms of the contract. The company reserves the right to seek an injunction to prevent employees from joining competitors. Employees are encouraged to consult with the legal department if they have questions about the scope of the non-compete clause.",
    "The contract also includes provisions for intellectual property (IP) ownership to protect the company’s interests. Any work, invention, or idea created by an employee during their employment belongs to the company. This includes software code, designs, reports, and marketing strategies developed on company time or using company resources. Employees must sign an IP assignment agreement as part of their onboarding process. The company retains the right to use, modify, or distribute employee-created IP without additional compensation. Employees are prohibited from using company IP for personal gain or sharing it with third parties. Violation of IP terms may result in termination and legal consequences. The company conducts regular training on IP policies to ensure employee awareness.",
    "Termination of employment, whether voluntary or involuntary, requires adherence to specific procedures outlined in the contract. Employees must provide 30 days’ written notice if resigning from their position. The company may waive the notice period at its discretion, with payment in lieu of notice. Upon termination, employees must return all company property, including laptops, access cards, and documents. A final paycheck will be issued within 7 business days, including any accrued benefits or unused leave pay. The company may conduct an exit interview to gather feedback and ensure a smooth transition. Employees are required to sign a termination agreement confirming compliance with all contract terms. Failure to follow termination procedures may impact the employee’s eligibility for rehire in the future.",
    "New employees are required to complete a 2-week onboarding program to familiarize themselves with company operations. The program includes training on company tools, such as the HR portal, email systems, and project management software. Employees will attend sessions on company policies, including attendance, leave, and workplace conduct. HR will provide a welcome kit containing a company handbook, ID badge, and branded merchandise. Each new employee will be assigned a mentor to guide them during their first 3 months. Mentors will schedule weekly check-ins to address questions and provide support. The onboarding program concludes with a feedback session to ensure new hires feel prepared for their roles.",
    "Full-time employees are eligible for a comprehensive benefits package starting on their first day of employment. Benefits include health insurance with coverage for medical, vision, and prescription drugs. Dental coverage is provided, including annual cleanings, fillings, and orthodontic care for dependents. Employees can enroll in a 401(k) retirement plan with a 5% company match on contributions. The company also offers life insurance and short-term disability coverage at no additional cost. Benefits enrollment occurs annually in November, with changes effective the following January. Employees must submit enrollment forms through the HR portal by the deadline. HR conducts benefits orientation sessions to help employees understand their options.",
    "Performance reviews are conducted biannually in June and December to evaluate employee contributions. Employees are assessed on key performance indicators (KPIs), including productivity, quality of work, and adherence to deadlines. Teamwork and collaboration are evaluated through peer feedback and manager observations. Employees will set SMART goals (Specific, Measurable, Achievable, Relevant, Time-bound) during each review cycle. Managers will provide constructive feedback and identify areas for improvement. High performers may be eligible for bonuses or promotions based on review outcomes. Employees who do not meet expectations will be placed on a performance improvement plan (PIP). The PIP includes specific targets and a 60-day timeline to demonstrate improvement.",
    "The company supports employee development through a variety of training and upskilling opportunities. Employees can access an online learning platform with courses on leadership, technical skills, and industry trends. The company reimburses up to $2,000 per year for external training programs or certifications. Employees must submit a reimbursement request with receipts within 30 days of completing the program. Monthly lunch-and-learn sessions are held to share knowledge on topics like time management and stress reduction. Employees are encouraged to attend at least one professional development event per quarter. Managers will work with employees to create individual development plans aligned with career goals. The HR team tracks participation in development programs to ensure equitable access for all staff.",
    "Employee wellness is a priority, and the company offers programs to support physical and mental health. A gym membership discount of 20% is available at partnered fitness centers nationwide. Employees can participate in weekly yoga and meditation sessions held virtually or in the office. The company provides an Employee Assistance Program (EAP) for confidential counseling on personal or work-related issues. Flu shots and health screenings are offered annually at no cost to employees. Standing desks and ergonomic chairs are available upon request to promote a healthy workspace. Employees are encouraged to take regular breaks to prevent burnout and maintain productivity. The HR team organizes quarterly wellness challenges, such as step-count competitions, with prizes for participants."
]

# Initialize the Hugging Face tokenizer
try:
    tokenizer = RagTokenizer.from_pretrained("facebook/rag-sequence-nq")
except Exception as e:
    print(f"Error loading tokenizer: {e}")
    raise

# Load the RAG model without the default retriever
try:
    model = RagSequenceForGeneration.from_pretrained("facebook/rag-sequence-nq")
except Exception as e:
    print(f"Error loading model: {e}")
    raise

# Simplified document retrieval for the specific query
def retrieve_documents(query, documents):
    query_lower = query.lower()
    # Directly retrieve the document for the specific query
    if "full time employees" in query_lower and "health insurance" in query_lower:
        # Document at index 11 is the relevant one
        retrieved_doc = documents[11]
        print(f"Retrieved document 11: {retrieved_doc[:100]}...")
        return retrieved_doc, "direct match"
    return None, None

# Function to generate an answer using the RAG model
def generate_answer(user_query):
    try:
        # Retrieve relevant documents
        retrieved_doc, score = retrieve_documents(user_query, documents)
        
        if retrieved_doc is None:
            return "Sorry, I couldn't find any relevant information in the documents to answer your query."

        # Log the retrieved document for debugging
        print(f"Retrieved document (score: {score}): {retrieved_doc}")

        # Tokenize the retrieved document and query
        input_text = f"{user_query} {retrieved_doc}"
        inputs = tokenizer(input_text, return_tensors="pt", truncation=True, max_length=512)

        # Generate an answer
        with torch.no_grad():
            outputs = model.generate(
                input_ids=inputs['input_ids'],
                attention_mask=inputs['attention_mask'],
                max_length=100,
                num_beams=5,
                early_stopping=True
            )

        # Decode and return the response
        answer = tokenizer.decode(outputs[0], skip_special_tokens=True)
        return answer
    except Exception as e:
        print(f"Error in generate_answer: {e}")
        return f"An error occurred: {str(e)}"

# Create a Gradio interface for testing the chatbot
iface = gr.Interface(fn=generate_answer, inputs="text", outputs="text", title="RAG Chatbot")

# Launch the interface
if __name__ == "__main__":
    iface.launch()