rollback

app.py

@@ -1,94 +1,63 @@
-# from fastapi import FastAPI, HTTPException
-# from pydantic import BaseModel
-# from transformers import AutoModelForCausalLM, AutoTokenizer
-# from typing import List
-# import torch
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from typing import List
+import torch
 
-# app = FastAPI(title="Language Model API")
+app = FastAPI(title="Language Model API")
 
-# # Model configuration
-# CHECKPOINT = "HuggingFaceTB/SmolLM2-135M-Instruct"
-# DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+# Model configuration
+CHECKPOINT = "HuggingFaceTB/SmolLM2-135M-Instruct"
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 
-# # Initialize model and tokenizer
-# try:
-#     tokenizer = AutoTokenizer.from_pretrained(CHECKPOINT)
-#     model = AutoModelForCausalLM.from_pretrained(CHECKPOINT).to(DEVICE)
-# except Exception as e:
-#     raise RuntimeError(f"Failed to load model: {str(e)}")
+# Initialize model and tokenizer
+try:
+    tokenizer = AutoTokenizer.from_pretrained(CHECKPOINT)
+    model = AutoModelForCausalLM.from_pretrained(CHECKPOINT).to(DEVICE)
+except Exception as e:
+    raise RuntimeError(f"Failed to load model: {str(e)}")
 
-# class ChatMessage(BaseModel):
-#     role: str
-#     content: str
+class ChatMessage(BaseModel):
+    role: str
+    content: str
 
-# class ChatRequest(BaseModel):
-#     messages: List[ChatMessage]
-#     max_new_tokens: int = 50
-#     temperature: float = 0.2
-#     top_p: float = 0.9
+class ChatRequest(BaseModel):
+    messages: List[ChatMessage]
+    max_new_tokens: int = 50
+    temperature: float = 0.2
+    top_p: float = 0.9
 
-# @app.post("/generate")
-# async def generate_response(request: ChatRequest):
-#     try:
-#         # Convert messages to the format expected by the model
-#         messages = [{"role": msg.role, "content": msg.content} for msg in request.messages]
+@app.post("/generate")
+async def generate_response(request: ChatRequest):
+    try:
+        # Convert messages to the format expected by the model
+        messages = [{"role": msg.role, "content": msg.content} for msg in request.messages]
         
-#         # Prepare input
-#         input_text = tokenizer.apply_chat_template(messages, tokenize=False)
-#         inputs = tokenizer.encode(input_text, return_tensors="pt").to(DEVICE)
+        # Prepare input
+        input_text = tokenizer.apply_chat_template(messages, tokenize=False)
+        inputs = tokenizer.encode(input_text, return_tensors="pt").to(DEVICE)
         
-#         # Generate response
-#         outputs = model.generate(
-#             inputs,
-#             max_new_tokens=request.max_new_tokens,
-#             temperature=request.temperature,
-#             top_p=request.top_p,
-#             do_sample=True
-#         )
+        # Generate response
+        outputs = model.generate(
+            inputs,
+            max_new_tokens=request.max_new_tokens,
+            temperature=request.temperature,
+            top_p=request.top_p,
+            do_sample=True
+        )
         
-#         # Decode and return response
-#         response_text = tokenizer.decode(outputs[0])
+        # Decode and return response
+        response_text = tokenizer.decode(outputs[0])
         
-#         return {
-#             "generated_text": response_text
-#         }
+        return {
+            "generated_text": response_text
+        }
     
-#     except Exception as e:
-#         raise HTTPException(status_code=500, detail=str(e))
-
-# if __name__ == "__main__":
-#     import uvicorn
-#     uvicorn.run(app, host="0.0.0.0", port=7860)
-
-
-from fastapi import FastAPI, HTTPException
-from pydantic import BaseModel
-from transformers import GPT2Tokenizer, GPT2LMHeadModel
-import torch
-
-# Initialize FastAPI app
-app = FastAPI()
-
-# Load GPT-2 model and tokenizer
-tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
-model = GPT2LMHeadModel.from_pretrained('gpt2')
-
-# Define input schema
-class TextInput(BaseModel):
-    text: str
-
-@app.post("/generate")
-async def process_text(input: TextInput):
-    try:
-        encoded_input = tokenizer(input.text, return_tensors='pt')
-        with torch.no_grad():
-            output = model.generate(**encoded_input, max_length=50, num_return_sequences=1)
-        generated_text = tokenizer.decode(output[0], skip_special_tokens=True)
-        return {"generated_text": generated_text}
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
 
-# Example root endpoint
-@app.get("/")
-async def root():
-    return {"message": "Welcome to the GPT-2 processing API!"}
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=7860)
+
+