Spaces:
Sleeping
Sleeping
| from fastapi import FastAPI | |
| from transformers import AutoModelForCausalLM, AutoTokenizer | |
| import torch | |
| MODEL_NAME = "deepseek-ai/deepseek-coder-6.7b-instruct" | |
| app = FastAPI() | |
| tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME) | |
| model = AutoModelForCausalLM.from_pretrained( | |
| MODEL_NAME, | |
| torch_dtype=torch.float32, | |
| device_map="cpu" | |
| ) | |
| def chat(req: dict): | |
| messages = req.get("messages", []) | |
| content = messages[-1]["content"] | |
| inputs = tokenizer(content, return_tensors="pt") | |
| outputs = model.generate( | |
| **inputs, | |
| max_new_tokens=256, | |
| temperature=0.7 | |
| ) | |
| response = tokenizer.decode(outputs[0], skip_special_tokens=True) | |
| return { | |
| "choices": [{ | |
| "message": {"role": "assistant", "content": response} | |
| }] | |
| } | |
| def root(): | |
| return {"status": "DeepSeek API is online"} | |