asd

app.py

@@ -1,7 +1,7 @@
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 import gradio as gr
-import spaces
+
 # Load the model and tokenizer
 model_name = "NoaiGPT/merged-llama3-8b-instruct-1720894657"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
@@ -11,15 +11,14 @@ model = AutoModelForCausalLM.from_pretrained(model_name)
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 model.to(device)
 
+# Create a text generation pipeline
+text_generator = pipeline("text-generation", model=model, tokenizer=tokenizer, device=0 if torch.cuda.is_available() else -1)
+
 # Define the prediction function
-@spaces.GPU
 def generate_text(prompt):
-    # Tokenize the input and move to GPU if available
-    inputs = tokenizer(prompt, return_tensors="pt").to(device)
-    # Generate text using the model
-    outputs = model.generate(inputs.input_ids, max_length=200, num_return_sequences=1)
-    # Decode the generated text
-    generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    # Generate text using the pipeline
+    outputs = text_generator(prompt, max_length=200, num_return_sequences=1)
+    generated_text = outputs[0]["generated_text"]
     return generated_text
 
 # Define the Gradio interface