Update README.md

README.md

@@ -29,39 +29,53 @@ pip install torch==2.0.0
 ```
 
 ```python
+from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
-from transformers import pipeline
 
-generate_text = pipeline(
-    model="diegomiranda/EleutherAI-70M-cypher-generator",
-    torch_dtype="auto",
-    trust_remote_code=True,
-    use_fast=True,
-    device_map={"": "cuda:0"},
-)
-
-res = generate_text(
-    "Why is drinking water so healthy?",
-    min_new_tokens=2,
-    max_new_tokens=500,
-    do_sample=False,
-    num_beams=2,
-    temperature=float(0.0),
-    repetition_penalty=float(1.0),
-    renormalize_logits=True
-)
-print(res[0]["generated_text"])
+def generate_response(prompt, model_name):
+    tokenizer = AutoTokenizer.from_pretrained(
+        model_name,
+        use_fast=True,
+        trust_remote_code=True,
+    )
+    
+    model = AutoModelForCausalLM.from_pretrained(
+        model_name,
+        torch_dtype=torch.float32,
+        device_map={"": "cpu"},
+        trust_remote_code=True,
+    )
+    model.cpu().eval()
+    
+    inputs = tokenizer(prompt, return_tensors="pt", add_special_tokens=False).to("cpu")
+    
+    tokens = model.generate(
+        input_ids=inputs["input_ids"],
+        attention_mask=inputs["attention_mask"],
+        min_new_tokens=2,
+        max_new_tokens=500,
+        do_sample=False,
+        num_beams=2,
+        temperature=float(0.0),
+        repetition_penalty=float(1.0),
+        renormalize_logits=True
+    )[0]
+    
+    tokens = tokens[inputs["input_ids"].shape[1]:]
+    answer = tokenizer.decode(tokens, skip_special_tokens=True)
+    
+    return answer
 ```
 
-You can print a sample prompt after the preprocessing step to see how it is feed to the tokenizer:
+# Example usage
 
 ```python
-print(generate_text.preprocess("Why is drinking water so healthy?")["prompt_text"])
+model_name = "diegomiranda/EleutherAI-70M-cypher-generator"
+prompt = "Create a Cypher statement to answer the following question:Retorne os processos de Direito Tributário que se baseiam em lei 939 de 1992?<|endoftext|>"
+response = generate_response(prompt, model_name)
+print(response)
 ```
 
-```bash
-Why is drinking water so healthy?<|endoftext|>
-```
 
 Alternatively, you can download [h2oai_pipeline.py](h2oai_pipeline.py), store it alongside your notebook, and construct the pipeline yourself from the loaded model and tokenizer. If the model and the tokenizer are fully supported in the `transformers` package, this will allow you to set `trust_remote_code=False`.