Testing Suggested Code Fix

app.py

@@ -1,36 +1,3 @@
-# import gradio as gr
-# from transformers import AutoTokenizer, AutoModelForCausalLM
-# import torch
-# import spaces
-
-# tokenizer = None
-# model = None
-
-# def loadmodel():
-#     tokenizer = AutoTokenizer.from_pretrained("ISTA-DASLab/Meta-Llama-3.1-70B-AQLM-PV-2Bit-1x16")
-#     model = AutoModelForCausalLM.from_pretrained("ISTA-DASLab/Meta-Llama-3.1-70B-AQLM-PV-2Bit-1x16", torch_dtype='auto', device_map='auto')
-#     return tokenizer, model
-
-# @spaces.GPU
-# def generate_text(prompt):
-#     global tokenizer, model
-#     if tokenizer is None or model is None:
-#         tokenizer, model = loadmodel()
-    
-#     inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
-#     outputs = model.generate(inputs.input_ids, max_length=100)
-#     return tokenizer.decode(outputs[0], skip_special_tokens=True)
-
-# interface = gr.Interface(
-#     fn=generate_text,
-#     inputs="text",
-#     outputs="text",
-#     title="Meta-Llama-3.1-70B Text Generation",
-#     description="Enter a prompt and generate text using Meta-Llama-3.1-70B.",
-# )
-
-# interface.launch()
-
 import spaces
 import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM
@@ -38,7 +5,7 @@ import torch
 import subprocess
 import os
 
-def install_cuda_toolkit():
+def install_cuda_toolkit(): #Swiftly Provided by https://huggingface.co/John6666 to fix OSError: CUDA_HOME environment variable is not set. Please set it to your CUDA install root.
     # CUDA_TOOLKIT_URL = "https://developer.download.nvidia.com/compute/cuda/11.8.0/local_installers/cuda_11.8.0_520.61.05_linux.run"
     CUDA_TOOLKIT_URL = "https://developer.download.nvidia.com/compute/cuda/12.2.0/local_installers/cuda_12.2.0_535.54.03_linux.run"
     CUDA_TOOLKIT_FILE = "/tmp/%s" % os.path.basename(CUDA_TOOLKIT_URL)
@@ -55,7 +22,7 @@ def install_cuda_toolkit():
     # Fix: arch_list[-1] += '+PTX'; IndexError: list index out of range
     os.environ["TORCH_CUDA_ARCH_LIST"] = "8.0;8.6"
 
-install_cuda_toolkit()
+install_cuda_toolkit() #Swiftly Provided by https://huggingface.co/John6666 to fix OSError: CUDA_HOME environment variable is not set. Please set it to your CUDA install root.
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
@@ -65,15 +32,15 @@ model = AutoModelForCausalLM.from_pretrained("ISTA-DASLab/Meta-Llama-3.1-70B-AQL
 @spaces.GPU
 def generate_text(prompt):
     inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
-    outputs = model.generate(inputs.input_ids, max_length=100)
+    outputs = model.generate(inputs.input_ids) #, max_length=100)
     return tokenizer.decode(outputs[0], skip_special_tokens=True)
 
 interface = gr.Interface(
     fn=generate_text,
     inputs="text",
     outputs="text",
-    title="Meta-Llama-3.1-70B Text Generation",
-    description="Enter a prompt and generate text using Meta-Llama-3.1-70B.",
+    title="Meta-Llama-3.1-70B-AQLM-PV-2Bit-1x16 Text Generation",
+    description="Enter a prompt and generate text using Meta-Llama-3.1-70B-AQLM-PV-2Bit-1x16. Responses are a little bit different Meta-Llama-3.1-70B",
 )
 
 interface.launch()