Update app.py

app.py

@@ -3,6 +3,12 @@ import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import time
 
+print(f"CUDA is available: {torch.cuda.is_available()}")
+print(f"CUDA device count: {torch.cuda.device_count()}")
+if torch.cuda.is_available():
+    print(f"Current CUDA device: {torch.cuda.current_device()}")
+    print(f"CUDA device name: {torch.cuda.get_device_name(0)}")
+
 class ConversationManager:
     def __init__(self):
         self.models = {}
@@ -24,7 +30,12 @@ class ConversationManager:
         try:
             print(f"Attempting to load model: {model_name}")
             tokenizer = AutoTokenizer.from_pretrained(model_name)
-            model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto", load_in_8bit=True)
+            try:
+                # Try to load the model with GPU support
+                model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto", load_in_8bit=True)
+            except RuntimeError as e:
+                print(f"GPU loading failed, falling back to CPU: {e}")
+                model = AutoModelForCausalLM.from_pretrained(model_name, device_map="cpu")
             self.models[model_name] = (model, tokenizer)
             print(f"Successfully loaded model: {model_name}")
             return self.models[model_name]
@@ -36,28 +47,14 @@ class ConversationManager:
 
     def generate_response(self, model_name, prompt):
         model, tokenizer = self.load_model(model_name)
-    
-        if "llama" in model_name.lower():
-            formatted_prompt = self.format_llama2_prompt(prompt)
-        else:
-            formatted_prompt = self.format_general_prompt(prompt)
-    
+        
+        formatted_prompt = f"Human: {prompt.strip()}\n\nAssistant:"
+        
         inputs = tokenizer(formatted_prompt, return_tensors="pt", max_length=1024, truncation=True)
         with torch.no_grad():
             outputs = model.generate(**inputs, max_length=200, num_return_sequences=1, do_sample=True)
         return tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-    def format_llama2_prompt(self, prompt):
-        B_INST, E_INST = "[INST]", "[/INST]"
-        B_SYS, E_SYS = "<<SYS>>\n", "\n<</SYS>>\n\n"
-        system_prompt = "You are a helpful AI assistant. Please provide a concise and relevant response."
-    
-        formatted_prompt = f"{B_INST} {B_SYS}{system_prompt}{E_SYS}{prompt.strip()} {E_INST}"
-        return formatted_prompt
-
-    def format_general_prompt(self, prompt):
-        return f"Human: {prompt.strip()}\n\nAssistant:"
-
     def add_to_conversation(self, model_name, response):
         self.conversation.append((model_name, response))
         if "task complete?" in response.lower():
@@ -181,20 +178,21 @@ def rewind_and_insert(steps, inserted_response, history):
     return manager.get_conversation_history(), ""
 
 open_source_models = [
-    "meta-llama/Llama-2-7b-chat-hf",
-    "meta-llama/Llama-2-13b-chat-hf",
-    "meta-llama/Llama-2-70b-chat-hf",
     "mistralai/Mixtral-8x7B-Instruct-v0.1",
     "bigcode/starcoder2-15b",
     "bigcode/starcoder2-3b",
     "tiiuae/falcon-7b",
-    "tiiuae/falcon-40b",
     "EleutherAI/gpt-neox-20b",
     "google/flan-ul2",
     "stabilityai/stablelm-zephyr-3b",
     "HuggingFaceH4/zephyr-7b-beta",
     "microsoft/phi-2",
-    "google/gemma-7b-it"
+    "google/gemma-7b-it",
+    "OpenAssistant/oasst-sft-4-pythia-12b-epoch-3.5",
+    "mosaicml/mpt-7b-chat",
+    "databricks/dolly-v2-12b",
+    "thebloke/Wizard-Vicuna-13B-Uncensored-HF",
+    "bigscience/bloom-560m"
 ]
 
 with gr.Blocks() as demo:
@@ -303,4 +301,4 @@ with gr.Blocks() as demo:
     delay_slider.change(lambda x: setattr(manager, 'delay', x), inputs=[delay_slider])
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()