Update app.py

app.py

@@ -1,64 +1,285 @@
+import os
+import re
+import json
 import gradio as gr
-from huggingface_hub import InferenceClient
+import requests
+from typing import List, Dict
+from googlesearch import search
+import google.generativeai as genai
+from google.generativeai.types import HarmCategory, HarmBlockThreshold
 
-"""
-For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
-"""
-client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")
+def initialize_gemini(api_key: str):
+    """Initialize the Google Gemini API with appropriate configurations"""
+    genai.configure(api_key=api_key)
+    generation_config = {
+        "temperature": 0.2,
+        "top_p": 0.8,
+        "top_k": 40,
+        "max_output_tokens": 1024,
+    }
+    safety_settings = {
+        HarmCategory.HARM_CATEGORY_HARASSMENT: HarmBlockThreshold.BLOCK_NONE,
+        HarmCategory.HARM_CATEGORY_HATE_SPEECH: HarmBlockThreshold.BLOCK_NONE,
+        HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT: HarmBlockThreshold.BLOCK_NONE,
+        HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT: HarmBlockThreshold.BLOCK_NONE,
+    }
+    
+    model = genai.GenerativeModel(
+        model_name="gemini-1.5-flash",
+        generation_config=generation_config,
+        safety_settings=safety_settings
+    )
+    return model
 
+def google_search_naics(company_name: str) -> List[str]:
+    """Find potential NAICS codes for a company using Google search"""
+    query = f"NAICS code 2022 for {company_name}"
+    naics_codes = set()
+    
+    try:
+        search_results = search(query, stop=5, pause=2)
+        
+        for result_url in search_results:
+            try:
+                response = requests.get(result_url, timeout=5)
+                if response.status_code == 200:
+                    # Extract 6-digit NAICS codes
+                    found_codes = re.findall(r'\b\d{6}\b', response.text)
+                    naics_codes.update(found_codes)
+            except Exception as e:
+                print(f"Error fetching {result_url}: {e}")
+        
+        return list(naics_codes)[:5]  # Return up to 5 extracted NAICS codes
+    except Exception as e:
+        print(f"Error performing Google search: {str(e)}")
+        return []
 
-def respond(
-    message,
-    history: list[tuple[str, str]],
-    system_message,
-    max_tokens,
-    temperature,
-    top_p,
-):
-    messages = [{"role": "system", "content": system_message}]
+def get_naics_classification(model, company_name: str, context: str, candidates: List[str]) -> dict:
+    """
+    Use Gemini AI to determine the most appropriate NAICS code from candidates
+    First provides reasoning, then multiple possibilities with confidence levels
+    """
+    try:
+        # If we have candidate codes from Google search
+        if candidates:
+            prompt = f"""
+You are a NAICS code classification expert. Based on the company information provided and the NAICS code candidates found from Google search, determine the most appropriate NAICS code.
 
-    for val in history:
-        if val[0]:
-            messages.append({"role": "user", "content": val[0]})
-        if val[1]:
-            messages.append({"role": "assistant", "content": val[1]})
+Company Name: {company_name}
+Context Information: {context}
 
-    messages.append({"role": "user", "content": message})
+NAICS Code Candidates from Google Search: {candidates}
 
-    response = ""
+First, explain your reasoning for which industry this company belongs to.
+Then list 3 potential NAICS classifications with confidence percentages (must add up to 100%).
+Finally, provide your final conclusion.
 
-    for message in client.chat_completion(
-        messages,
-        max_tokens=max_tokens,
-        stream=True,
-        temperature=temperature,
-        top_p=top_p,
-    ):
-        token = message.choices[0].delta.content
-
-        response += token
-        yield response
+Your response should be in this format:
+REASONING: [Your detailed reasoning about the company's industry classification]
 
+POSSIBILITY_1: [Industry name] - NAICS Code [6-digit code] - [XX]% confidence
+POSSIBILITY_2: [Industry name] - NAICS Code [6-digit code] - [XX]% confidence  
+POSSIBILITY_3: [Industry name] - NAICS Code [6-digit code] - [XX]% confidence
 
+CONCLUSION: I am [XX]% confident this company is [industry description] which is NAICS code [6-digit code]
 """
-For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
+        # If no candidates were found from Google search
+        else:
+            prompt = f"""
+You are a NAICS code classification expert. Based on the company information provided, determine the most appropriate NAICS code.
+
+Company Name: {company_name}
+Context Information: {context}
+
+First, explain your reasoning for which industry this company belongs to.
+Then list 3 potential NAICS classifications with confidence percentages (must add up to 100%).
+Finally, provide your final conclusion.
+
+Your response should be in this format:
+REASONING: [Your detailed reasoning about the company's industry classification]
+
+POSSIBILITY_1: [Industry name] - NAICS Code [6-digit code] - [XX]% confidence
+POSSIBILITY_2: [Industry name] - NAICS Code [6-digit code] - [XX]% confidence  
+POSSIBILITY_3: [Industry name] - NAICS Code [6-digit code] - [XX]% confidence
+
+CONCLUSION: I am [XX]% confident this company is [industry description] which is NAICS code [6-digit code]
 """
-demo = gr.ChatInterface(
-    respond,
-    additional_inputs=[
-        gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
-        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
-        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
-        gr.Slider(
-            minimum=0.1,
-            maximum=1.0,
-            value=0.95,
-            step=0.05,
-            label="Top-p (nucleus sampling)",
-        ),
-    ],
-)
+        response = model.generate_content(prompt)
+        response_text = response.text.strip()
+        
+        # Extract reasoning
+        reasoning_match = re.search(r'REASONING:(.*?)POSSIBILITY_1:', response_text, re.DOTALL | re.IGNORECASE)
+        reasoning = reasoning_match.group(1).strip() if reasoning_match else "No reasoning provided."
+        
+        # Extract possibilities
+        possibilities = []
+        
+        # Try to extract possibility 1
+        poss1_match = re.search(r'POSSIBILITY_1:(.*?)POSSIBILITY_2:', response_text, re.DOTALL | re.IGNORECASE)
+        if poss1_match:
+            possibilities.append(poss1_match.group(1).strip())
+            
+        # Try to extract possibility 2  
+        poss2_match = re.search(r'POSSIBILITY_2:(.*?)POSSIBILITY_3:', response_text, re.DOTALL | re.IGNORECASE)
+        if poss2_match:
+            possibilities.append(poss2_match.group(1).strip())
+            
+        # Try to extract possibility 3
+        poss3_match = re.search(r'POSSIBILITY_3:(.*?)CONCLUSION:', response_text, re.DOTALL | re.IGNORECASE)
+        if poss3_match:
+            possibilities.append(poss3_match.group(1).strip())
+        
+        # Extract conclusion
+        conclusion_match = re.search(r'CONCLUSION:(.*?)
+    except Exception as e:
+        print(f"Error getting NAICS classification: {str(e)}")
+        return {
+            "naics_code": "000000",
+            "reasoning": f"Error analyzing company: {str(e)}"
+        }
+
+def find_naics_code(api_key, company_name, company_description):
+    """Main function to find NAICS code that will be called by Gradio"""
+    if not api_key or not company_name:
+        return "Please provide both API key and company name."
+    
+    try:
+        # Initialize Gemini API
+        model = initialize_gemini(api_key)
+        
+        # Search for NAICS candidates
+        naics_candidates = google_search_naics(company_name)
+        
+        # Get classification
+        if not naics_candidates:
+            result = get_naics_classification(model, company_name, company_description, [])
+        else:
+            result = get_naics_classification(model, company_name, company_description, naics_candidates)
+        
+        # Format the output
+        output = f"## NAICS Code for {company_name}\n\n"
+        output += f"**NAICS Code:** {result['naics_code']}\n\n"
+        output += f"**Reasoning:**\n{result['reasoning']}\n\n"
+        
+        # Add possibilities section
+        if 'possibilities' in result and result['possibilities']:
+            output += f"**Possible Classifications:**\n\n"
+            for i, possibility in enumerate(result['possibilities'], 1):
+                output += f"{i}. {possibility}\n\n"
+        
+        # Add conclusion
+        if 'conclusion' in result and result['conclusion']:
+            output += f"**Conclusion:**\n{result['conclusion']}\n\n"
+        
+        if naics_candidates:
+            output += f"**Candidate NAICS Codes Found from Google:**\n{', '.join(naics_candidates)}"
+        
+        return output
+    
+    except Exception as e:
+        return f"Error: {str(e)}"
+
+# Create Gradio Interface
+with gr.Blocks(title="NAICS Code Finder") as app:
+    gr.Markdown("# NAICS Code Finder")
+    gr.Markdown("This app helps you find the appropriate NAICS code for a company based on its name and description.")
+    
+    with gr.Row():
+        with gr.Column():
+            api_key = gr.Textbox(label="Google Gemini API Key", placeholder="Enter your Gemini API key here", type="password")
+            company_name = gr.Textbox(label="Company Name", placeholder="Enter the company name")
+            company_description = gr.Textbox(label="Company Description", placeholder="Enter a brief description of the company", lines=5)
+            
+            submit_btn = gr.Button("Find NAICS Code")
+        
+        with gr.Column():
+            output = gr.Markdown(label="Result")
+    
+    submit_btn.click(
+        fn=find_naics_code,
+        inputs=[api_key, company_name, company_description],
+        outputs=output
+    )
+
+if __name__ == "__main__":
+    app.launch()
+, response_text, re.DOTALL | re.IGNORECASE)
+        conclusion = conclusion_match.group(1).strip() if conclusion_match else "No conclusion provided."
+        
+        # Extract final NAICS code from conclusion
+        naics_match = re.search(r'NAICS code (\d{6})', conclusion)
+        if naics_match:
+            naics_code = naics_match.group(1)
+        else:
+            # Try to find any 6-digit code in the conclusion
+            code_match = re.search(r'\b(\d{6})\b', conclusion)
+            naics_code = code_match.group(1) if code_match else "000000"
+        
+        return {
+            "naics_code": naics_code,
+            "reasoning": reasoning,
+            "possibilities": possibilities,
+            "conclusion": conclusion
+        }
+    except Exception as e:
+        print(f"Error getting NAICS classification: {str(e)}")
+        return {
+            "naics_code": "000000",
+            "reasoning": f"Error analyzing company: {str(e)}"
+        }
+
+def find_naics_code(api_key, company_name, company_description):
+    """Main function to find NAICS code that will be called by Gradio"""
+    if not api_key or not company_name:
+        return "Please provide both API key and company name."
+    
+    try:
+        # Initialize Gemini API
+        model = initialize_gemini(api_key)
+        
+        # Search for NAICS candidates
+        naics_candidates = google_search_naics(company_name)
+        
+        # Get classification
+        if not naics_candidates:
+            result = get_naics_classification(model, company_name, company_description, [])
+        else:
+            result = get_naics_classification(model, company_name, company_description, naics_candidates)
+        
+        # Format the output
+        output = f"## NAICS Code for {company_name}\n\n"
+        output += f"**NAICS Code:** {result['naics_code']}\n\n"
+        output += f"**Reasoning:**\n{result['reasoning']}\n\n"
+        
+        if naics_candidates:
+            output += f"**Candidate NAICS Codes Found:**\n{', '.join(naics_candidates)}"
+        
+        return output
+    
+    except Exception as e:
+        return f"Error: {str(e)}"
 
+# Create Gradio Interface
+with gr.Blocks(title="NAICS Code Finder") as app:
+    gr.Markdown("# NAICS Code Finder")
+    gr.Markdown("This app helps you find the appropriate NAICS code for a company based on its name and description.")
+    
+    with gr.Row():
+        with gr.Column():
+            api_key = gr.Textbox(label="Google Gemini API Key", placeholder="Enter your Gemini API key here", type="password")
+            company_name = gr.Textbox(label="Company Name", placeholder="Enter the company name")
+            company_description = gr.Textbox(label="Company Description", placeholder="Enter a brief description of the company", lines=5)
+            
+            submit_btn = gr.Button("Find NAICS Code")
+        
+        with gr.Column():
+            output = gr.Markdown(label="Result")
+    
+    submit_btn.click(
+        fn=find_naics_code,
+        inputs=[api_key, company_name, company_description],
+        outputs=output
+    )
 
 if __name__ == "__main__":
-    demo.launch()
+    app.launch()