Use distillgpt

main.py

@@ -6,49 +6,49 @@ from transformers import pipeline
 from transformers import AutoTokenizer, AutoModelForCausalLM
 
 
-# Load pre-trained tokenizer and model (Works)
-# model_name = "distilgpt2"
-# tokenizer = AutoTokenizer.from_pretrained(model_name)
-# model = AutoModelForCausalLM.from_pretrained(model_name)
+Load pre-trained tokenizer and model (Works)
+model_name = "distilgpt2"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(model_name)
 
-# # Example usage: Generate text
-# prompt = "The quick brown fox"
-# input_ids = tokenizer.encode(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)
+# Example usage: Generate text
+prompt = "The quick brown fox"
+input_ids = tokenizer.encode(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)
 
-# print(generated_text)
+print(generated_text)
 
 
-import transformers
-import torch
-import logging
+# import transformers
+# import torch
+# import logging
 
-model_id = "deepcogito/cogito-v1-preview-llama-3B"
+# 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",
-)
+# pipeline = transformers.pipeline(
+#     "text-generation",
+#     model=model_id,
+#     model_kwargs={"torch_dtype": torch.bfloat16},
+#     device_map="auto",
+# )
 
 
-print("Pipeline loaded")
-logging.info("Pipeline loaded")
+# 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."},
-]
+# 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,
-)
+# outputs = pipeline(
+#     messages,
+#     max_new_tokens=512,
+# )
 
-logging.info("Generated text")
-print(outputs[0]["generated_text"][-1])
+# logging.info("Generated text")
+# print(outputs[0]["generated_text"][-1])
 
 
 app = FastAPI()
@@ -62,21 +62,27 @@ class Item(BaseModel):
 
 @app.post("/generate/")
 async def generate_text(item: Item):
-    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."},
-    ]
+    # 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,
-    )
+    # outputs = pipeline(
+    #     messages,
+    #     max_new_tokens=512,
+    # )
 
-    logging.info("request got")
+    # logging.info("request got")
 
-    resp = outputs[0]["generated_text"][-1]
+    # resp = outputs[0]["generated_text"][-1]
 
-    logging.info("Response generated")
+    # 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)
+
+    resp = generated_text
 
     return {"response": resp}