Update main.py

main.py

@@ -9,7 +9,12 @@ from transformers import AutoTokenizer, AutoModelForCausalLM
 #Load pre-trained tokenizer and model (Works)
 model_name = "gpt2"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(
+    model_name,
+    device_map="auto",
+    torch_dtype="auto"
+)
+
 
 # Example usage: Generate text
 prompt = "The quick brown fox"
@@ -20,37 +25,6 @@ generated_text = tokenizer.decode(output[0], skip_special_tokens=True)
 print(generated_text)
 
 
-# import transformers
-# import torch
-# import logging
-
-# model_id = "deepcogito/cogito-v1-preview-llama-3B"
-
-# pipeline = transformers.pipeline(
-#     "text-generation",
-#     model=model_id,
-#     model_kwargs={"torch_dtype": torch.bfloat16},
-#     device_map="auto",
-# )
-
-
-# print("Pipeline loaded")
-# logging.info("Pipeline loaded")
-
-# messages = [
-#     {"role": "system", "content": "You are a pirate chatbot who always responds in pirate speak!"},
-#     {"role": "user", "content": "Give me a short introduction to LLMs."},
-# ]
-
-# outputs = pipeline(
-#     messages,
-#     max_new_tokens=512,
-# )
-
-# logging.info("Generated text")
-# print(outputs[0]["generated_text"][-1])
-
-
 app = FastAPI()
 
 class EchoMessage(BaseModel):
@@ -78,11 +52,22 @@ async def generate_text(item: Item):
 
     # logging.info("Response generated")
 
-    input_ids = tokenizer.encode(item.prompt, return_tensors="pt")
-    output = model.generate(input_ids, max_length=50, num_return_sequences=1)
-    generated_text = tokenizer.decode(output[0], skip_special_tokens=True)
+    inputs = tokenizer(prompt, return_tensors="pt", padding=True, return_attention_mask=True).to(model.device)
+
+
+    # input_ids = tokenizer.encode(item.prompt, return_tensors="pt")
+    # output = model.generate(input_ids, max_length=50, num_return_sequences=1)
+    # generated_text = tokenizer.decode(output[0], skip_special_tokens=True)
+    # resp = generated_text
+
+    outputs = model.generate(
+        **inputs,
+        max_new_tokens=100,
+        pad_token_id=tokenizer.eos_token_id  # Set this to suppress warning
+    )
+
+    resp = tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-    resp = generated_text
 
     return {"response": resp}