Use coito

Dockerfile

@@ -6,6 +6,14 @@ COPY ./requirements.txt /code/requirements.txt
 
 RUN pip install --no-cache-dir --upgrade -r /code/requirements.txt
 
+RUN pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu
+
+
+# sudo apt-get install lzma
+# sudo apt-get install liblzma-dev
+# sudo apt-get install libbz2-dev
+
+
 COPY . .
 
 CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "7860"]

README.md

@@ -8,3 +8,17 @@ pinned: false
 ---
 
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+
+## Instructions for local use:
+
+1. Build with docker
+
+```
+docker build -t app:v1 .
+```
+
+2. Run
+
+```
+docker run -p 7860:7860 app:v1
+```

main.py

@@ -3,8 +3,34 @@ from pydantic import BaseModel
 from huggingface_hub import InferenceClient
 import uvicorn
 from transformers import pipeline
+from transformers import AutoTokenizer, AutoModelForCausalLM
 
-model_name = "distilgpt2"
+
+# 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)
+
+# print(generated_text)
+
+
+import transformers
+import torch
+
+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",
+)
 
 
 app = FastAPI()
@@ -16,29 +42,20 @@ class Item(BaseModel):
     prompt: str
 
 
-def generate(item: Item):
-    generator = pipeline("text-generation", model=model_name)
-
-    # Your input prompt
-    prompt = item.prompt
+@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."},
+    ]
 
-    # Generate text
-    generated_texts = generator(
-        prompt,
-        max_length=50,  # Maximum length of the generated text
-        num_return_sequences=1,  # Number of different sequences to generate
-        temperature=0.8,  # Controls the randomness of the output
-        top_k=50,        # Limits the number of top tokens to consider
-        top_p=0.95,      # Nucleus sampling parameter
-        do_sample=True   # Enable sampling for non-deterministic output
+    outputs = pipeline(
+        messages,
+        max_new_tokens=512,
     )
 
-    return generated_texts
-
-
-@app.post("/generate/")
-async def generate_text(item: Item):
-    return {"response": generate(item)}
+    resp = outputs[0]["generated_text"][-1]
+    return {"response": resp}
 
 @app.get("/")
 async def home():

requirements.txt

@@ -2,4 +2,5 @@ fastapi
 uvicorn
 huggingface_hub
 pydantic
-transformers
+transformers
+accelerate