Update app.py

app.py

@@ -1,15 +1,15 @@
 import os
 import subprocess
 import streamlit as st
-from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
+from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer, AutoConfig, AutoModel, pipeline_utils
 import black
 from pylint import lint
 from io import StringIO
-import openai
 import sys
+from huggingface_hub import notebook_login
 
-# Set your OpenAI API key here
-openai.api_key = "YOUR_OPENAI_API_KEY"
+notebook_login()
+hf_token = "YOUR_HF_TOKEN"
 
 HUGGING_FACE_REPO_URL = "https://huggingface.co/spaces/acecalisto3/DevToolKit"
 PROJECT_ROOT = "projects"
@@ -41,7 +41,6 @@ class AIAgent:
         agent_prompt = f"""
 As an elite expert developer, my name is {self.name}. I possess a comprehensive understanding of the following areas:
 {skills_str}
-
 I am confident that I can leverage my expertise to assist you in developing and deploying cutting-edge web applications. Please feel free to ask any questions or present any challenges you may encounter.
 """
         return agent_prompt
@@ -68,6 +67,9 @@ I am confident that I can leverage my expertise to assist you in developing and
 
         return summary, next_step
 
+def load_hf_token():
+    return hf_token
+
 def save_agent_to_file(agent):
     """Saves the agent's prompt to a file."""
     if not os.path.exists(AGENT_DIRECTORY):
@@ -102,7 +104,9 @@ def chat_interface_with_agent(input_text, agent_name):
     try:
         model = AutoModelForCausalLM.from_pretrained(model_name)
         tokenizer = AutoTokenizer.from_pretrained(model_name)
-        generator = pipeline("text-generation", model=model, tokenizer=tokenizer)
+        generator = pipeline("text-generation", model=model, tokenizer=tokenizer,
+                             framework="pt",
+                             model_kwargs={'load_in_8bit': True, 'torch_dtype': torch.float16})
     except EnvironmentError as e:
         return f"Error loading model: {e}"
 
@@ -113,7 +117,7 @@ def chat_interface_with_agent(input_text, agent_name):
     if input_ids.shape[1] > max_input_length:
         input_ids = input_ids[:, :max_input_length]
 
-    outputs = model.generate(
+    outputs = generator(
         input_ids, max_new_tokens=50, num_return_sequences=1, do_sample=True,
         pad_token_id=tokenizer.eos_token_id  # Set pad_token_id to eos_token_id
     )