Update app.py

app.py

@@ -13,21 +13,16 @@ examplecofee = """A cortado is a Spanish beverage consisting of espresso mixed w
 
 model_path = "NousResearch/Genstruct-7B"
 
-# hf_token = os.getenv("HF_TOKEN")
-# if not hf_token:
-#     raise ValueError("Hugging Face token not found. Please set the HF_TOKEN environment variable.")
-
 tokenizer = AutoTokenizer.from_pretrained(model_path)
-# quantization_config = BitsAndBytesConfig(load_in_8bit=True)
-model = AutoModelForCausalLM.from_pretrained( model_path, device_map='cuda', load_in_8bit=True)
+quantization_config = BitsAndBytesConfig(load_in_8bit=True)
+model = AutoModelForCausalLM.from_pretrained( model_path, device_map='cuda', quantization_config=quantization_config)
 
 @spaces.GPU
 def generate_text(prompt, custom_prompt, temperature, max_length):
-    structured_input = f"[Title]: {prompt}\n[Content]: {custom_prompt}"
-    inputs = tokenizer(structured_input, return_tensors='pt').to('cuda')    
-    generated_ids = model.generate(**inputs, max_new_tokens=max_length, temperature=temperature, do_sample=True).split(tokenizer.eos_token)[0])
+    structured_input = f"[[[Title]]]: {prompt}\n[[[Content]]]: {custom_prompt}"
+    inputs = tokenizer(structured_input, return_tensors='pt').to('cuda')
+    generated_ids = model.generate(**inputs, max_new_tokens=max_length, temperature=temperature, pad_token_id=tokenizer.eos_token_id)
     generated_text = tokenizer.decode(generated_ids[0], skip_special_tokens=True)
-    
     return generated_text