Update README.md

README.md

@@ -34,23 +34,35 @@ This repository includes the weights learned during the training process. It sho
 
 <!-- This section is for the model use without fine-tuning or plugging into a larger ecosystem/app. -->
 ```python
-from transformers import AutoModelForSequenceClassification, AutoTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 
 # Load the tokenizer, adjust configuration if needed
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForCausalLM.from_pretrained(model_name)
 
-# Load the fine-tuned model with its trained weights
-fine_tuned_model = AutoModelForSequenceClassification.from_pretrained(
-    'fionazhang/mistral_7b_environment',
+# Text generation
+def generate_text_sequences(pipe, prompt):
+    sequences = pipe(
+        f"prompt",
+        do_sample=True,
+        max_new_tokens=100,
+        temperature=0.8,
+        top_k=50,
+        top_p=0.95,
+        num_return_sequences=1,
+    )
+    return sequences[0]['generated_text']
+
+# Now you can use the model for inference
+pipe = pipeline(
+    "text-generation",
+    model=model,
+    tokenizer=tokenizer,
+    torch_dtype=torch.bfloat16,
+    device_map="auto",
+    pad_token_id=2
 )
-
-# Now you can use `fine_tuned_model` for inference or further training
-input_text = "The impact of climate change on"
-output_text = fine_tuned_model.generate(tokenizer.encode(input_text, return_tensors="pt"))
-
-print(tokenizer.decode(output_text[0], skip_special_tokens=True))
-
+print(generate_text_sequences(pipe, "your prompt"))
 ```