Update app.py

app.py

@@ -1,28 +1,14 @@
-import transformers
+from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
-import os
-from huggingface_hub import login
 
-read_key = os.environ.get("Read_Token")
-login(token=read_key)
-
-model_id = "meta-llama/Meta-Llama-3.1-8B-Instruct"
-
-pipeline = transformers.pipeline(
-    "text-generation",
-    model=model_id,
-    model_kwargs={"torch_dtype": torch.bfloat16},
-    device_map="auto",
+tokenizer = AutoTokenizer.from_pretrained("google/gemma-2b-it")
+model = AutoModelForCausalLM.from_pretrained(
+    "google/gemma-2b-it",
+    torch_dtype=torch.bfloat16
 )
 
-messages = [
-    {"role": "system", "content": "You are a pirate chatbot who always responds in pirate speak!"},
-    {"role": "user", "content": "Who are you?"},
-]
+input_text = "Write me a poem about Machine Learning."
+input_ids = tokenizer(input_text, return_tensors="pt")
 
-outputs = pipeline(
-    messages,
-    max_new_tokens=256,
-    pad_token_id = pipeline.tokenizer.eos_token_id
-)
-print(outputs[0]["generated_text"][-1])
+outputs = model.generate(**input_ids)
+print(tokenizer.decode(outputs[0]))