Update app.py

app.py

@@ -209,6 +209,8 @@ gemini = os.getenv("GEMINI")
 genai.configure(api_key=gemini)
 model = genai.GenerativeModel("gemini-1.5-flash")
 
+
+
 def read_project_files(project_path):
     """Reads all files in the project directory and its subdirectories."""
     file_paths = []
@@ -270,17 +272,13 @@ def identify_required_functions(project_path, functionality_description):
     return response.text
 
 
-# Load the Qwen model and tokenizer
-model_name = "Qwen/Qwen2.5-Coder-7B-Instruct"
-hf_model = AutoModelForCausalLM.from_pretrained(
-    model_name,
-    torch_dtype="auto",
-    device_map="auto"
-)
-tokenizer = AutoTokenizer.from_pretrained(model_name)
+# Hugging Face Inference API endpoint for the model
+API_URL = "https://api-inference.huggingface.co/models/Qwen/Qwen2.5-Coder-32B-Instruct"
+qwen = os.getenv("QWEN")
+headers = {"Authorization": f"Bearer {qwen}"}
 
-def validate_and_generate_documentation(hf_model, tokenizer, gemini_output, file_contents, functionality_description):
-    """Uses Qwen model to validate functions and generate full documentation."""
+def validate_and_generate_documentation(api_url, headers, gemini_output, file_contents, functionality_description):
+    """Uses the Hugging Face Inference API to validate functions and generate documentation."""
     # Generate the prompt for the Qwen model
     prompt = f"""
     User-specified functionality: '{functionality_description}'
@@ -290,7 +288,7 @@ def validate_and_generate_documentation(hf_model, tokenizer, gemini_output, file
     Project files:
     """
     for file_path, content in file_contents.items():
-        # Truncate content to avoid exceeding token limits
+        # Truncate content to avoid excessively large payloads
         truncated_content = content[:1000] if len(content) > 1000 else content
         prompt += f"File: {os.path.basename(file_path)}\n{truncated_content}\n\n"
 
@@ -318,30 +316,15 @@ def validate_and_generate_documentation(hf_model, tokenizer, gemini_output, file
       - Data structures: <Details of data structures used>
     """
 
-    # Prepare the chat-style input for Qwen
-    messages = [
-        {"role": "system", "content": "You are Qwen, created by Alibaba Cloud. You are a helpful assistant."},
-        {"role": "user", "content": prompt}
-    ]
-    text = tokenizer.apply_chat_template(
-        messages,
-        tokenize=False,
-        add_generation_prompt=True
-    )
-    model_inputs = tokenizer([text], return_tensors="pt", truncation=True, max_length=32768).to(hf_model.device)
-
-    # Generate output from the model
-    generated_ids = hf_model.generate(
-        **model_inputs,
-        max_new_tokens=2048
-    )
-    generated_ids = [
-        output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
-    ]
-
-    # Decode and return the response
-    response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
-    return response
+    # Send the prompt to the Hugging Face API
+    payload = {"inputs": prompt}
+    response = requests.post(api_url, headers=headers, json=payload)
+
+    # Handle the API response
+    if response.status_code == 200:
+        return response.json().get("generated_text", "No output generated.")
+    else:
+        raise ValueError(f"Error: {response.status_code}, {response.text}")
 
 def generate_documentation_page():
     st.subheader(f"Generate Documentation for {st.session_state.current_project}")
@@ -368,9 +351,9 @@ def generate_documentation_page():
                     file_paths = read_project_files(project_folder)
                     file_contents = read_files(file_paths)
                     
-                    # Call Hugging Face model for validation and documentation
+                    # Call the Hugging Face API for validation and documentation
                     final_documentation = validate_and_generate_documentation(
-                        hf_model, tokenizer, gemini_result, file_contents, functionality
+                        API_URL, headers, gemini_result, file_contents, functionality
                     )
                     
                     # Display the final documentation