Update app.py

app.py

@@ -5,6 +5,8 @@ import os
 import google.generativeai as genai
 import zipfile
 from git import Repo
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
 
 
 # Database setup
@@ -267,6 +269,63 @@ def identify_required_functions(project_path, functionality_description):
     # Process and return the response
     return response.text
 
+
+# Load Hugging Face model and tokenizer
+tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-Coder-32B-Instruct")
+hf_model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2.5-Coder-32B-Instruct", device_map="auto")
+
+def validate_and_generate_documentation(hf_model, tokenizer, gemini_output, file_contents, functionality_description):
+    """Uses Hugging Face model to validate functions and generate full documentation."""
+    # Generate the prompt for the Hugging Face model
+    prompt = f"""
+    The user-provided functionality: '{functionality_description}'
+    The functions identified by Gemini:
+    {gemini_output}
+
+    Project files:
+    """
+    for file_path, content in file_contents.items():
+        prompt += f"File: {os.path.basename(file_path)}\n{content[:1000]}...\n\n"  # Truncate large files for the model
+
+    prompt += """
+    Task:
+    1. Validate if the functions identified by Gemini are sufficient for implementing the functionality.
+    2. If not, identify all additional functions required.
+    3. For all relevant functions, generate detailed documentation in the following format:
+    """
+    prompt += """
+    Project Summary:
+    Summary of the entire project, making sure to mention the language it's programmed in and any libraries or other dependencies it has
+
+    Functionality Summary:
+    Summary of the user-specified functionality
+
+    Functionality Flow:
+    How the programmer goes from inputting information into the first function to the last function and its output to complete
+    the functionality that is described by the user. Make sure to mention each function that is used, and how inputs and outputs flow between each other.
+
+    Function Documentation:
+    For each file that contains the relevant functions:
+      For each function determined as relevant within the current file:
+        Summary:
+          summarize what the function does
+        Inputs:
+          the inputs and their data types, and their relevance in the scope of the specified functionality
+        Outputs:
+          the output, its data type, and its relevance in the scope of the specified functionality
+        Dependencies:
+          the dependencies of the function and where they come from
+        Data structures:
+          the data structures that the function relies on
+    """
+    # Encode and call the Hugging Face model
+    inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=2048).to(hf_model.device)
+    outputs = hf_model.generate(inputs["input_ids"], max_length=4096, num_return_sequences=1)
+
+    # Decode the response
+    decoded_output = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return decoded_output
+
 def generate_documentation_page():
     st.subheader(f"Generate Documentation for {st.session_state.current_project}")
     st.write("Enter the functionality or parts of the project for which you'd like to identify relevant functions.")
@@ -283,18 +342,23 @@ def generate_documentation_page():
             user_folder = os.path.join("user_projects", st.session_state.username)
             project_folder = os.path.join(user_folder, st.session_state.current_project)
             
-            # Ensure compatibility with GitHub repositories
             if os.path.exists(project_folder):
                 try:
-                    # Call the function to identify relevant functions
-                    result = identify_required_functions(project_folder, functionality)
-
-                    # Check if Gemini returned an error message
-                    if "No relevant functions found" in result:
-                        st.error("No relevant functions were found for the specified functionality. Please refine your input.")
-                    else:
-                        st.success("Relevant functions identified successfully!")
-                        st.text_area("Relevant Functions", result, height=400)
+                    # Call Gemini to identify required functions
+                    gemini_result = identify_required_functions(project_folder, functionality)
+                    
+                    # Read project files
+                    file_paths = read_project_files(project_folder)
+                    file_contents = read_files(file_paths)
+                    
+                    # Call Hugging Face model for validation and documentation
+                    final_documentation = validate_and_generate_documentation(
+                        hf_model, tokenizer, gemini_result, file_contents, functionality
+                    )
+                    
+                    # Display the final documentation
+                    st.success("Documentation generated successfully!")
+                    st.text_area("Generated Documentation", final_documentation, height=600)
                 except Exception as e:
                     st.error(f"An error occurred: {e}")
             else:
@@ -309,6 +373,7 @@ def generate_documentation_page():
 
 
 
+
 #------------------------------------------------------------------------------------------------------------------------------------------------------------------------